Barrier domain doesn't match implemented domain

Hey all, I know this feature (overridden_get_grid_sizes_for_insn_ids) is technically unsupported, but I imagine there are other ways to set something like this up such that there is a mismatch between a desired and implemented domain. In either case, I need to find a way to get this working (and could use some pointers in the right direction)

MWE: https://gist.github.com/arghdos/854fc7d2058824f96cc30a6f3b631844

The problem (I think): A barrier inside a kernel "extended" by overridden_get_grid_sizes_for_insn_ids inherits the codegen_state's implemented domain (loopy/codegen/control.py:165). The implemented domain for the assignment instructions a and b then pull the updated implemented domain from the insn_inames (loopy/codegen/instruction.py:94) and then the barrier's implemented domain no longer matches the desired

How to fix: Suggestions most welcome 😕