Verifier and Type inference changes for reduction operations #1869
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
from
279b326
to
c668b0e
Compare
|
We need a means of checking if an IR downgrade is possible in terms of constraints. Let me think about how this should happen. Perhaps we can have "CompatibilityConstraintChecks" somewhere. Other option would be to add ReduceOpV2 in VHLO, but for constraint modification we may be able to do better than this. |
ghpvnist
reviewed
Dec 5, 2023
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
2 times, most recently
from
476495d
to
41d5788
Compare
ghpvnist
approved these changes
Dec 13, 2023
|
Put together main...GleasonK:stablehlo:reduce-compat Please port those changes over and implement constraint versioning for these changed ops. Aside from that, add compatibility tests per: https://github.com/openxla/stablehlo/blob/main/docs/vhlo.md#add--update-unit-tests - including positive and negative tests. Can see CollectiveBroadcastOp for example: #1856 |
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
3 times, most recently
from
d3d72c3
to
cd14e1a
Compare
|
@GleasonK Updated the PR with the suggested changes. PTAL. |
sdasgup3
commented
Dec 15, 2023
GleasonK
approved these changes
Dec 16, 2023
GleasonK
requested changes
Dec 16, 2023
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
from
cd14e1a
to
fae94af
Compare
GleasonK
requested changes
Dec 20, 2023
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
from
fae94af
to
b284312
Compare
GleasonK
approved these changes
Dec 20, 2023
sdasgup3
force-pushed
the
reduce-verifier-impl
branch
from
b284312
to
e497deb
Compare
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 12, 2024
…le semantics [RFC](openxla/stablehlo#1664), to MHLO. The CL does two things: 1. The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase. 2. There exists canonicalization passes like `group-reduction-dimensions` and `hlo-canonicalize-reduction` which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region [example](https://github.com/openxla/xla/blob/a91877b9c9aa1edf307c5927782111b1a81cd81d/xla/mlir_hlo/mhlo/transforms/group_reduction_dimensions/group_reduction_dimensions.cc#L228). This is problematic as, with the [change](openxla/stablehlo#1869), the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 597407271
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 12, 2024
…le semantics [RFC](openxla/stablehlo#1664), to MHLO. The CL does two things: 1. The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase. 2. There exists canonicalization passes like `group-reduction-dimensions` and `hlo-canonicalize-reduction` which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region [example](https://github.com/openxla/xla/blob/a91877b9c9aa1edf307c5927782111b1a81cd81d/xla/mlir_hlo/mhlo/transforms/group_reduction_dimensions/group_reduction_dimensions.cc#L228). This is problematic as, with the [change](openxla/stablehlo#1869), the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 597407271
sdasgup3
added a commit
that referenced
this pull request
Jan 19, 2024
The upstream change #1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions in [hlo_ops.cc](https://github.com/openxla/xla/blob/main/xla/mlir_hlo/mhlo/IR/hlo_ops.cc) uses those APIs. The PR updates the visibility and signature of those API for a clearer integration. Specifically, the PR does the followings: 1. **updates `getAccumulatorTypes` to return a error status when the input regions is empty**: This function is used in type inference of various reduction based operations ([eg](https://github.com/openxla/stablehlo/blob/d5b464925371092095ac934b46ba93ebd4284223/stablehlo/dialect/TypeInference.cpp#L2589)). This functions enables infering type based on the reduction block of the operation, which is something proposed in [RFC](#1664). However, there could be [instances](https://github.com/openxla/xla/blob/a91877b9c9aa1edf307c5927782111b1a81cd81d/xla/mlir_hlo/mhlo/transforms/group_reduction_dimensions/group_reduction_dimensions.cc#L228) when type inference can be called with empty region in which case we would like to report a meaningful diagnostic message. 2. **Allow `hlo::inferAllReduceOp` to accept multiple operands information**: In stableHLO, `all_reduce` op have a single operand ([e.g.](https://github.com/openxla/stablehlo/blob/d5b464925371092095ac934b46ba93ebd4284223/stablehlo/dialect/StablehloOps.td#L1355)), whereas in MHLO the op can take multiple operand ([e.g.](https://github.com/openxla/xla/blob/79aba0801ef75c1c2dffbb4ecc506a0d8144c9ac/xla/mlir_hlo/mhlo/IR/hlo_ops.td#L1528). The `hlo::inferAllReduceOp` signature is updated to accommodate both cases. 3. Remove unused arguments to functions `verifyReduceOpInputsAndInferShape` and `inferReduceOp`.
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599684600
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599684600
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599684600
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599684600
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599684600
copybara-service bot
pushed a commit
to tensorflow/mlir-hlo
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599930190
copybara-service bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599930190
copybara-service bot
pushed a commit
to openxla/xla
that referenced
this pull request
Jan 19, 2024
Other than the usual integration, the CL does two things: The upstream change openxla/stablehlo#1869 in StableHLO updates various API related to shape inference. MHLO shape inference functions uses those APIs. The CL fixes the invocation of those APIs in MHLO codebase so as to sync the semantics of StableHLO reduction operation with MHLO. There exists canonicalization passes like group-reduction-dimensions and hlo-canonicalize-reduction which create reduce operation using builder methods that calls type inference of reduce op with empty reduction region example. This is problematic as, with the change, the type inference of reduce op is now dependent on the reduction body. The CL updates all the calls sites of the problematic builder (the one which calls type inference with empty reduction block) with the invocation of a new custom builder method introduced for mhlo::Reduce operation. Note that at the moment we do not need similar custom builder for other reduction based operations (like scatter, reduce_scatter, all_reduce, select_and_scatter, reduce_window) as they are presently created using a builder version take result type as an input and hence does not call inference from within. Also, the CL adds verification tests for the operations with promotable semantics. PiperOrigin-RevId: 599930190
This was referenced Jan 29, 2024
sdasgup3
added a commit
that referenced
this pull request
Jan 30, 2024
The PR implements the a custom `reduce` op builder similar to what we have for mhlo [code](https://github.com/openxla/xla/blob/50aec2b3b54ce7a861f45bc3b0ae9b2cc2ee2a28/xla/mlir_hlo/mhlo/IR/hlo_ops.cc#L3917). ## Background #1869 allows the block arguments of reduce op to have different element types than that of the input arguments of reduce op and the output element type of the reduce op has to equal to those block arguments. As a consequence the output type of reduce op can no longer be inferred from the operand types. The auto-generated builders creates a reduce op with empty block and, as a result, does not allow inferring the type. The proposed solution is to create a custom builder which takes the element-type of the block arguments as arguments allowing result type inference.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Implements the specification changes at #1796.
The PR adds/updates the verifier and type inference routines for the following ops:
reduce, reduce_window, select_and_scatter, all_reduce, reduce_scatter, scatter. Please refer to #1796 for the updated constraints which the PR implements. Note the #1796 is going to be merged soon.Here are the changes for each operation:
reduce
reduceand other related ops #1796 added a new constraint C8reduce_window:
reduceand other related ops #1796 updated the C16select_and_scatter
reduceand other related ops #1796 added a new constraint C12scatter
- Specification of quantized
reduceand other related ops #1796 added a new constraint C17- Updated labels
- Add positive tests ; negative tests verifying scatter_c15 at verify_scatter.mlir and type inference tests at infrer_stablehlo.mlir, for scatter_C17.
reduce_scatter
reduceand other related ops #1796 added a new constraint C9SameOperandsAndResultElementTypeimplementing the outdated constraint. For the new constraint C9, we added a check atverifyReduceScatterOp.all_reduce
reduceand other related ops #1796 added a new constraint C7inferReturnTypeComponentsFromOperands. For the new constraint C9, we added a traitInferTensorTypein the tablegen definition of the op.