Actually… no, I don’t need to keep them in the same IRModule. Having 2 different IRModules and then compiling then separately would work perfectly for me.
Do you think its easier to partition the IRModule and get 2 different IRModules?
EDIT: I took a look at the partition_conversions function, because I think I could do something similar to obtain the 2 different IRModules. I thought I could use it out-of-the-box, but then noticed my use case is a little different:
As far as I understand, this is what the partition_conversions function expects as a graph:
And it creates 3 IRModules, which are then combined into one with 4 functions in it:
- A quantize function, which calls the quantize_op
- A quantized_main function, which calls all the quantized op_int_*
- A dequantize function, which calls only the dequantize_op
- A main function which calls the 3 functions in the correct order, as if the partition_conversions transformation was not executed.
My use case is different because my graph looks something like this (of course, this is a simplified version):
/ / \ \
out_1 out_2 out_3 out_4
Basically what I would like to achieve is to partition this into 2 (or 3, if we also partition the first quantize_op in a different IRModule) different modules:
- One containing all ops until both dequantize_ops (so this would be an IRModule with 2 outputs)
- One containing all the ops afterwards (so 2 input and 4 outputs)
I believe the partition_conversions function is breaking because I have ops after the dequantize_ops, and also because I have 2 of them.