Relax AlterOpLayout with AXIS_SEPARATOR

Questions
3

Thank you @psrivas2 for your response! Sorry for coming back late on this but I was going through the source code to understand the pass better. I’m planning to add support for AXIS_SEPARATOR and have proposed below mentioned changes after going through the code. Can you please review them and share your feedback?

  1. To update inputs with index_map, TransformLayout is invoked here. I plan to add AXIS_SEPARATOR in LayoutTransformAttrs to be saved for respective buffers. We can get the AXIS_SEPARATOR from user in the same way we are getting index_map.
  2. Similarly for output buffers, TransformLayout gets invoked in TransformLayoutInverse function and AXIS_SEPARATOR support can be added in the same manner as point 1 above.
  3. The replacement_primfunc passed to AlterOpImpl pass by user will already have support for AXIS_SEPARATOR because it’s a user configurable function.

Please let me know if any point above needs any further clarification and I will be happy to answer. As always, thank you for your help! :slight_smile: