search

openxla / stablehlo

Backward compatible ML compute opset inspired by HLO/MHLO
Apache License 2.0
357 stars 100 forks source link

Lower shape.cstr_broadcastable op in ShapeLegalizeToHLO #2384

Closed sdasgup3 closed 3 weeks ago

sdasgup3 commented 4 weeks ago

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 https://github.com/openxla/stablehlo/pull/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

sdasgup3 commented 3 weeks ago

GitHub issue #2390 focuses on exploring the removal of chlo-broadcast operation dependencies from stablehlo-legalize-quant-to-int pass. If this removal is successful, we can avoid further changes to shape legalization, as suggested in the pull request. Consequently, let's close the pull request for now.