Lower shape.cstr_broadcastable op in ShapeLegalizeToHLO #2384
Conversation
| // expected-error@+1 {{failed to legalize operation 'shape.cstr_broadcastable' that was explicitly marked illegal}} | ||
| %0 = shape.cstr_broadcastable %arg0, %arg1 : !shape.shape, !shape.shape | ||
| shape.assuming %0 { | ||
| } |
There was a problem hiding this comment.
Is this needed? Same comment below, I feel like this was covered in the first test.
| if (this->legalize_constraints_) { | ||
| target->addLegalOp<shape::ConstWitnessOp, shape::AssumingOp, | ||
| shape::AssumingYieldOp>(); | ||
| } |
There was a problem hiding this comment.
IMO it would be cleaner to consolidate the handling of this->legalize_constraints_ into a single if. If it's true, add the pattern to the set of patterns and add the legal ops.
There was a problem hiding this comment.
I agree. The current choice is driven by how the code is currenty structured: addition of legal ops handled separately from adding the pattern. Using a single if would create an outlier for legalize_contriants_ case. wdyt?
There was a problem hiding this comment.
Right, I think both takes are valid, but IMO in this case it would be more readable to have a single if since that would make it clearer what happens when legalize_constraints_ is true. Up to you though.
sdasgup3
force-pushed
the
support-shape-cstr-broadcastable
branch
from
c6d6244
to
7a44cab
Compare
There was a problem hiding this comment.
This is a blocker comment for merging this PR.
It seems that the current use-case of lowering statically shaped quantized program to stablehlo math operations can be achieve even without this pass.
With that it might be the case that we do not prioritize this PR for now.
I will get back with the resolution of this PR asap.
@mlevesquedion In anycase thanks for the review. They are exptremly valuable.
|
GitHub issue #2390 focuses on exploring the removal of chlo-broadcast operation dependencies from |
Lower shape.cstr_broadcastable op in ShapeLegalizeToHLO
Lower the op to shape_assertion custom_calls. And replace the result shape.witness with a const_witness of true. This allows the subsequent shape.assuming regions to be removed.
Also add an option to the pass to control whether to lower cstr_xxx ops. This makes the impl easier as we can reuse util functions in the file. And it makes sure existing use cases are not modified.
Why needed
The evaluation of quantized program depends bout lowering the stablehlo quantized types/ops to their integer equivalents. This is accomplished by pass like #2383 which introduces chlo boradcast operations. These chlo operations needs to be lowered to stablehlo ops as well using a pass which the current PR upgrades. Without this upgrade shape.cstr_broadcastable will not be lowered to stablehlo for evaluation.
Inspired from https://github.com/openxla/xla/blob/25f38dd29b8442e45072fc097d0f1ebe3163bf9d/xla/mlir_hlo/mhlo/transforms/mhlo_passes.td#L361