Skip to content

Commit

Permalink
[Stablehlo] fix compareOp with scalar's lowering (#3518)
Browse files Browse the repository at this point in the history
* use lhs tensor's element type as compute type when rhs is scalar.
* previously `a != 1.0`(a is a fp32 tensor) will lowering to `%6 =
stablehlo.compare EQ, %4, %5, FLOAT : (tensor<2x5xf64>, tensor<2x5xf64>)
-> tensor<2x5xi1>`
* now it will lowering to `%6 = stablehlo.compare EQ, %4, %5, FLOAT :
(tensor<2x5xf32>, tensor<2x5xf32>) -> tensor<2x5xi1>`
  • Loading branch information
qingyunqu authored Jul 2, 2024
1 parent e2fbded commit f1e3701
Showing 1 changed file with 2 additions and 0 deletions.
2 changes: 2 additions & 0 deletions lib/Conversion/TorchToStablehlo/Basic.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -517,6 +517,8 @@ class ConvertAtenCompareOp : public OpConversionPattern<AtenOpT> {
if (!rhsTy) {
rhs = hlo::scalarToStablehloTensor(rewriter, op, adaptor.getOther(),
rhs.getType());
// use lhs's element type as compute type
rhs = hlo::promoteType(rewriter, op.getLoc(), rhs, lhsTy);
rhsTy = dyn_cast<RankedTensorType>(rhs.getType());
}

Expand Down

0 comments on commit f1e3701

Please sign in to comment.