search

iree-org / iree

A retargetable MLIR-based machine learning compiler and runtime toolkit.
http://iree.dev/
Apache License 2.0
2.84k stars 612 forks source link

tfl.while and tfl.if legalization to tosa #9863

Open scottcantreid opened 2 years ago

scottcantreid commented 2 years ago

I will be writing a pass to convert tfl.while/tfl.if to tosa.while/tosa.if, respectively.

I'm taking inspiration from the implementation of the TOSA->SCF conversion. When completed, this pass should also lower all tfl.yield operations to tosa.yield.

As an external contributor, I will contribute this pass locally to iree_tf_compiler/TFL. Ultimately, this probably belongs within the tensorflow repository. @rsuderman has kindly offered to help migrate this into tensorflow.

allieculp commented 2 years ago

Thanks, Scott for filing this! @rsuderman Please help add priority to this issue.

scottcantreid commented 2 years ago

I believe I have a working pass for while, but I want to clarify the signature for if before I write that one.

In the tfl dialect, IfOp has only a single argument: TFL_BoolTensor:$cond

However, in the tosa dialect, IfOp has the an equivalent boolean argument, as well as a Variadic<Tosa_Tensor>:$inputs. Does anyone know the purpose of this variadic argument?

jpienaar commented 2 years ago

I'd expect this to be used as form before export given spec form https://www.mlplatform.org/tosa/tosa_spec.html#_cond_if, but ops aren't isolated from above so shouldn't be needed in general AFAICS

scottcantreid commented 2 years ago

Thanks! So if I understand correctly, a tensor needs to be passed into input_list in order to be in scope in the then/else regions. This is different from scf.if and tfl.if, where any variable previously defined in the graph is in scope.

scottcantreid commented 2 years ago

but ops aren't isolated from above so shouldn't be needed in general AFAICS

Given this, I can pretty easily write a lowering that leaves the input_list empty. I wonder if the TOSA group would be open to changing the spec to remove this unnecessary argument.

jpienaar commented 2 years ago

Yes, so I think I'd do what you propose and then we could have (like we do for TFLite) a "prepare for export" pass that makes it args. But for all optimizations, I'd prefer it in region form. And j don't think for iree it matters, only place that seems like it could matter is for export to interpreter/TOSA flatbuffer