Split verifiers and shape functions for ops with regions

[zhouxin913]

This is a pure refactor PR which fixes #400 (deprecated) and #623 (new)

As in the #623 , "we only need this for ops with regions", a full list of ops with region extract from https://github.com/openxla/stablehlo/blob/main/stablehlo/dialect/StablehloOps.td includes 11 ops in total:

Op	what's done
AllReduceOp	No change (already split)
CaseOp	No change (region is indispensable for type inference)
IfOp	No change (region is indispensable for type inference)
MapOp	No change (region is indispensable for type inference)
ReduceOp	Split
ReduceScatterOp	No change (Type Inference implementation on hold see #725)
ReduceWindowOp	Split
ScatterOp	No change (already split)
SelectAndScatterOp	No change (already split)
SortOp	Split
WhileOp	Split

The ideal split is that verifiers contain as almost all verifications, and the shape functions are simple as possible, but note:

1. IfOp/CaseOp/MapOp: We need info from region(s) to infer the return type, so an init function without regions is always invalid and should not exist. No change for them in this PR.
2. As both verifier & shape function need verification of the inputs/attrs, so we need put them in a separate utils functions.
   ReduceOp: introduce new util verifyReduceOpInputsAndInferShape()
ReduceWindow: introduce new util verifyReduceWindowOpInputsAndInferWindow()
   In each op, verifier does (1) call this new util function (2) verify region
   shape function: (1) call this new util function (2) generate inferred type from the intermediate result from (1)
3. Besides, the verification logic needs further fix see Verification of ReduceOp and ReduceWindowOp is not strict enough. #394, but this is out of scope of this PR.

[GleasonK]

Is the purpose of unused arguments just to make the integration easier? There's a good chance that this change breaks the integ because clang-tidy will (should?) fail on unused arguments. If we did want to go the unused route, would suggest the [[maybe_unused]] attribute.

Personal preference is to remove unused variables. We own both clients of the shared infrastructure so I don't see too big of a reason to prefer stability.

[zhouxin913]

Is the purpose of unused arguments just to make the integration easier? There's a good chance that this change breaks the integ because clang-tidy will (should?) fail on unused arguments. If we did want to go the unused route, would suggest the [[maybe_unused]] attribute.

Personal preference is to remove unused variables. We own both clients of the shared infrastructure so I don't see too big of a reason to prefer stability.

Thanks for the reminder! @GleasonK Yes I decide to go without necessary args. If in some case the an arg is not used, I will go with arg type but without arg name: to avoid clang-tidy error in both StableHLO/MHLO.

[burmako]

Excellent PR description! I agree with both the analysis and the implementation.

What are our next steps for CaseOp, IfOp and MapOp? You've made a convincing case that regions are indispensable for type inference of these ops, but in that case what do we do with the autogenerated builders? These questions don't have to be answered in this PR, but we should have a ticket about this, and I think that this ticket should have high priority.

[zhouxin913]

Excellent PR description! I agree with both the analysis and the implementation.

What are our next steps for CaseOp, IfOp and MapOp? You've made a convincing case that regions are indispensable for type inference of these ops, but in that case what do we do with the autogenerated builders? These questions don't have to be answered in this PR, but we should have a ticket about this, and I think that this ticket should have high priority.

Thanks for the quick review!
Opened #738 and let us think and track in this new bug: whether/how to fix this.

[smit-hinsu]

by the time we reach verifyReduceOp, we would have already inferred and checked that the results are compatible so this will result in duplicated call.

we may still need to calculate result dimensions as the op result may have incomplete (but still valid) shape. we could do that directly.

Documentation has a section on (complicated) verification flow.
https://github.com/llvm/llvm-project/blob/main/mlir/docs/DefiningDialects/Operations.md#custom-verifier-code

Separate from this change, another thing I realized is that we need to move verifyReducerShape to verifyRegions as some things like region's blocks have terminator op are verified later on. (We could also add HasOnlyGraphRegion trait to our ops with regions.)

[zhouxin913]

Let me defense with these:

1. As far as digging to code, verify() runs before inferReturnTypeComponents(), I have shared you with a doc of my founds.
2. This PR aims to address another tricky case that: the shape function is called alone without the verifier: thus we have no choice but to duplicate these verification logic. See Split verifiers and shape functions for ops with regions #623

Thanks for pointing out the final part in the comment! Could you open a bug for it and elaborate a bit? I feel it's an optimization, and could do it when we have more cycles. Please correct me otherwise.

[smit-hinsu]

1. As far as digging to code, verify() runs before inferReturnTypeComponents(), I have shared you with a doc of my founds.

you are right. I didn't realize that InferTypeOpInterface sets verifyWithRegions to 1.

That means that we need to move all the verification to verifyRegions by using hasRegionVerifier. By the time it reaches verifyRegions, inference has been verified so we don't have to do those checks again.