Issue using MLIR for tensorrt!

similar to introduced TensorRT in Tensorflow. I need to transform TF Dialect operations to TensorRT Dialect operations. I want to use DRR to achieve it. I found the ‘patternrewriter’ generated by DRR contains two steps: match and rewrite. The two steps is in the one function.
the ‘match’ step used to judge whether a TF Dialect operation can be converted to be a TensorRT Dialect operation.
the ‘rewrite’ step convert the TF Dialect operation to the TensorRT Dialect operation.

In fact, we only convert the TF Dialect operation to the TensorRT Dialect operation when the subgraph size is more than a specified number。 when we pick up the tensorrt subgraph, we need the “match step” result. So, I want to know is there a way to only generate the “match” step code?
best wishes to you!

I would do this slightly different than I think you have in mind. You can take advantage of MLIR’s ability to nest IR. First do a analysis pass where you segment the graph into regions. You may want to enforce a minimum on the TRT segment sizes. Create an op type which can encapsulate the TRT segments. After the segmentation, move segments of the graph into your “TRT region op”. This would act like a staging area where you can now construct the TRT graph and build the engine. You can think of this as “outlining code to a function” and look at how that’s accomplished elsewhere. You might have a other ops that you use after this to represent loading and calling the engine, etc. A feature of the “tf2tensorrt” compiler in TF is that you can retain the original segments in case your conversion fails. This would be a useful feature to have as well.

1 Like

I see, thank you very much