[BOO] Remove force_single_dispatch

[IanWood1]

Needs iree-org/iree@35a872c to ensure all conv kernels get fused to a single dispatch. Specifically, the pads need to be fused into the dispatches.

[IanWood1]

Keeping this as a draft until IREE is updated

[zjgarvey]

How do we know if this is working correctly? Can we add some tests to verify this doesn't regress, etc?

[ftynse]

Can we rather plumb this through as a flag (defaulted to False if needed). There are cases for layer norm that are currently failing to distribute without forcing single dispatch, in particular when intermediate results are forwarded from forward to reverse mode since those results don't participate in the "main" reduction. See https://discord.com/channels/689900678990135345/689932666505920596/1412414159248691210 for some more context.

[zjgarvey]

Can we rather plumb this through as a flag (defaulted to False if needed). There are cases for layer norm that are currently failing to distribute without forcing single dispatch, in particular when intermediate results are forwarded from forward to reverse mode since those results don't participate in the "main" reduction. See https://discord.com/channels/689900678990135345/689932666505920596/1412414159248691210 for some more context.

Yeah, I just noticed this too. We could maybe put it into the OpSignature class, or allow the layer norm regressions and try to address it in the compiler? @MaheshRavishankar do you have preferences on how we handle this?

[MaheshRavishankar]

Can we rather plumb this through as a flag (defaulted to False if needed). There are cases for layer norm that are currently failing to distribute without forcing single dispatch, in particular when intermediate results are forwarded from forward to reverse mode since those results don't participate in the "main" reduction. See https://discord.com/channels/689900678990135345/689932666505920596/1412414159248691210 for some more context.

We will have to fix that forward. Could you file an issue for what you are seeing?

[ftynse]

Why the rush? Can't we just treat this as an optimization option?

[MaheshRavishankar]

Why the rush? Can't we just treat this as an optimization option?

We can keep it as an optimization option, but it would be better to keep it default off. The default on is breaking a lot of cases. The issue that you pointed to is triaged and I think there is a fix for that already.

[IanWood1]

@ftynse do you have any ideas/preference on how this would be exposed as an optimization option?

[IanWood1]

@ftynse do you have any ideas/preference on how this would be exposed as an optimization option?

Nevermind I think the commit that @zjgarvey added should give you what you want.

[zjgarvey]

Yeah, the change I made should work as a temporary workaround until something like #1134 lands (where the fusion spec for an op already determines whether it uses the attr).

[zjgarvey]

LGTM