Add more MHLO folders to StableHLO (#2753)

Some of these may seem unsafe (simplifying float `X*1.0` or `X*0.0` may
not account for NaNs properly), but both MHLO and XLA simplifications do
this optim so probably generally safe.

Author: GleasonK

stablehlo/tests/transforms/stablehlo_aggressive_folder.mlir

@@ -269,6 +269,17 @@ func.func @eval_convert_f64_precision_loss() -> (tensor<1xf32>, tensor<f32>) {
 
 // -----
 
+// CHECK-LABEL: func @fold_sqrt
+func.func @fold_sqrt() -> (tensor<f32>) {
+  // CHECK: [[RESULT0:%.*]] = stablehlo.constant dense<2.0{{.*}}> : tensor<f32>
+  // CHECK: return [[RESULT0]]
+  %0 = stablehlo.constant dense<4.0> : tensor<f32>
+  %1 = stablehlo.sqrt %0 : tensor<f32>
+  func.return %1 : tensor<f32>
+}
+
+// -----
+
 // CHECK-LABEL: func @eval_transpose
 func.func @eval_transpose() -> (tensor<2x3x2xi32>, tensor<2x4x3xi32>, tensor<4x3x2xi32>) {
   // CHECK: [[RESULT0:%.*]] = stablehlo.constant dense<

stablehlo/tests/transforms/stablehlo_aggressive_simplification.mlir

@@ -899,6 +899,23 @@ func.func @multiply_by_one(%arg0: tensor<i32>) -> tensor<i32> {
   return %0 : tensor<i32>
 }
 
+// CHECK-LABEL: @multiply_by_zero_float
+func.func @multiply_by_zero_float(%arg0: tensor<f32>) -> tensor<f32> {
+  %cst = stablehlo.constant dense<0.0> : tensor<f32>
+  // CHECK: stablehlo.constant dense<0.0{{.*}}> : tensor<f32>
+  %0 = stablehlo.multiply %cst, %arg0 : tensor<f32>
+  return %0 : tensor<f32>
+}
+
+// CHECK-LABEL: @multiply_by_one_float
+func.func @multiply_by_one_float(%arg0: tensor<f32>) -> tensor<f32> {
+  %cst = stablehlo.constant dense<1.0> : tensor<f32>
+  %0 = stablehlo.multiply %cst, %arg0 : tensor<f32>
+  // CHECK-NOT: stablehlo.constant
+  // CHECK: return %arg0 : tensor<f32>
+  return %0 : tensor<f32>
+}
+
 // -----
 
 /////////

stablehlo/transforms/optimization/StablehloAggressiveFolder.cpp

@@ -12,9 +12,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
+#include <cmath>
 #include <cstddef>
 #include <cstdint>
+#include <functional>
 #include <numeric>
+#include <optional>
 #include <utility>
 
 #include "llvm/ADT/APInt.h"
@@ -23,11 +26,14 @@ limitations under the License.
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/LogicalResult.h"
 #include "mlir/Dialect/CommonFolders.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
+#include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
@@ -86,7 +92,7 @@ APSInt getAPSInt(Type type, uint64_t value) {
       /*isUnsigned=*/isUnsigned);
 }
 
-LogicalResult validateResultTypeForEval(PatternRewriter& rewriter,
+LogicalResult validateStaticShapeResult(PatternRewriter& rewriter,
                                         Operation* op, ShapedType resultType) {
   if (!resultType.hasStaticShape())
     return rewriter.notifyMatchFailure(
@@ -212,7 +218,7 @@ template <typename OpType, typename FuncType>
 LogicalResult evalElementwise(PatternRewriter& rewriter, OpType op,
                               FuncType fn) {
   auto resultType = op.getType();
-  if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+  if (failed(validateStaticShapeResult(rewriter, op, resultType)))
     return failure();
 
   if (!isa<IntegerType>(resultType.getElementType()))
@@ -279,6 +285,28 @@ struct FoldAddOpPattern final : OpRewritePattern<mlir::stablehlo::AddOp> {
   }
 };
 
+// A base class to use for patterns that may be used for integer shape math,
+// but also may be used for general folding of floats.
+template <typename OpType>
+struct ShapeOpRewritePattern : public OpRewritePattern<OpType> {
+  ShapeOpRewritePattern(MLIRContext* context, PatternBenefit benefit,
+                        bool foldFloat_)
+      : OpRewritePattern<OpType>(context, benefit), foldFloat{foldFloat_} {}
+
+  using OpRewritePattern<OpType>::OpRewritePattern;
+  using OpRewritePattern<OpType>::matchAndRewrite;
+
+  LogicalResult validateShapeFoldDtype(PatternRewriter& rewriter, OpType op,
+                                       ShapedType resultType) const {
+    if (resultType.getElementType().isInteger()) return success();
+    if (foldFloat && isa<FloatType>(resultType.getElementType()))
+      return success();
+    return rewriter.notifyMatchFailure(op, "skipping fold of shape op dtype");
+  }
+
+  bool foldFloat;
+};
+
 struct EvalAddOpShapePattern : public OpRewritePattern<AddOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(AddOp op,
@@ -327,7 +355,7 @@ struct EvalBroadcastInDimOpPattern : public OpRewritePattern<BroadcastInDimOp> {
   LogicalResult matchAndRewrite(BroadcastInDimOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     auto operandType = op.getOperand().getType();
@@ -442,7 +470,7 @@ struct EvalConcatenateOpPattern : public OpRewritePattern<ConcatenateOp> {
   LogicalResult matchAndRewrite(ConcatenateOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     if (op.getDimension() != 0)
@@ -460,31 +488,23 @@ struct EvalConcatenateOpPattern : public OpRewritePattern<ConcatenateOp> {
   }
 };
 
-struct EvalConvertOpPattern : public OpRewritePattern<ConvertOp> {
-  using OpRewritePattern::OpRewritePattern;
-
-  EvalConvertOpPattern(MLIRContext* context, PatternBenefit benefit,
-                       bool foldFloat_)
-      : OpRewritePattern<ConvertOp>(context, benefit), foldFloat{foldFloat_} {}
+struct EvalConvertOpPattern : public ShapeOpRewritePattern<ConvertOp> {
+  using ShapeOpRewritePattern::ShapeOpRewritePattern;
 
   LogicalResult matchAndRewrite(ConvertOp op,
                                 PatternRewriter& rewriter) const override {
     auto operand = op.getOperand();
     RankedTensorType resultType = op.getType();
 
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)) ||
+        failed(validateShapeFoldDtype(rewriter, op, resultType)))
       return failure();
 
     auto operandElemType = getElementTypeOrSelf(operand.getType());
     auto resultElemType = getElementTypeOrSelf(resultType);
-    if (!(operandElemType.isInteger() && resultElemType.isInteger()) &&
-        !foldFloat)
-      return rewriter.notifyMatchFailure(op,
-                                         "lossy computations are not allowed");
-
-    if (!resultElemType.isIntOrFloat())
-      return rewriter.notifyMatchFailure(
-          op, "expected integer or float result tensor type");
+    if (!foldFloat &&
+        (isa<FloatType>(operandElemType) || isa<FloatType>(resultElemType)))
+      return rewriter.notifyMatchFailure(op, "skipping fold of float convert");
 
     DenseIntOrFPElementsAttr elements;
     if (!matchPattern(operand, m_Constant(&elements)))
@@ -493,9 +513,6 @@ struct EvalConvertOpPattern : public OpRewritePattern<ConvertOp> {
 
     return evalConvert(rewriter, op, elements, resultType);
   }
-
- private:
-  bool foldFloat;
 };
 
 struct EvalDivOpPattern : public OpRewritePattern<DivOp> {
@@ -513,7 +530,7 @@ struct EvalGetDimensionSizeOpPattern
   LogicalResult matchAndRewrite(GetDimensionSizeOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     auto operandType = op.getOperand().getType();
@@ -552,23 +569,11 @@ struct FoldMulOpPattern final : OpRewritePattern<mlir::stablehlo::MulOp> {
 
   LogicalResult matchAndRewrite(mlir::stablehlo::MulOp op,
                                 PatternRewriter& rewriter) const override {
-    auto elemType = op.getType().getElementType();
-    Value lhs = op.getLhs();
-    Value rhs = op.getRhs();
-
     TypedAttr lhsAttr;
-    matchPattern(lhs, m_Constant(&lhsAttr));
+    matchPattern(op.getLhs(), m_Constant(&lhsAttr));
 
     TypedAttr rhsAttr;
-    matchPattern(rhs, m_Constant(&rhsAttr));
-
-    // The canonical form has the constant operand as the RHS.
-    if (isa<IntegerType>(elemType) && lhsAttr && !rhsAttr) {
-      rewriter.modifyOpInPlace(op, [op, lhs, rhs] {
-        op->setOperands(ValueRange{rhs, lhs});
-      });
-      return success();
-    }
+    matchPattern(op.getRhs(), m_Constant(&rhsAttr));
 
     if (TypedAttr res;
         lhsAttr && rhsAttr &&
@@ -613,16 +618,18 @@ struct EvalRemOpPattern : public OpRewritePattern<RemOp> {
   }
 };
 
-struct EvalReshapeOpPattern : public OpRewritePattern<ReshapeOp> {
-  using OpRewritePattern::OpRewritePattern;
+struct EvalReshapeOpPattern : public ShapeOpRewritePattern<ReshapeOp> {
+  using ShapeOpRewritePattern::ShapeOpRewritePattern;
+
   LogicalResult matchAndRewrite(ReshapeOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)) ||
+        failed(validateShapeFoldDtype(rewriter, op, resultType)))
       return failure();
 
     // Pattern: reshape(cst, shape) -> cst
-    DenseIntElementsAttr attr;
+    DenseIntOrFPElementsAttr attr;
     if (!matchPattern(op.getOperand(), m_Constant(&attr)))
       return rewriter.notifyMatchFailure(op, "expected constant operand");
     rewriter.replaceOpWithNewOp<ConstantOp>(op, attr.reshape(resultType));
@@ -635,7 +642,7 @@ struct EvalSelectOpPattern : public OpRewritePattern<SelectOp> {
   LogicalResult matchAndRewrite(SelectOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     SmallVector<APSInt> pred, onTrue, onFalse;
@@ -717,7 +724,7 @@ struct EvalSliceOpPattern : public OpRewritePattern<SliceOp> {
   LogicalResult matchAndRewrite(SliceOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     auto operand = op.getOperand();
@@ -778,6 +785,37 @@ struct EvalSubtractOpPattern : public OpRewritePattern<SubtractOp> {
   }
 };
 
+struct FoldSqrtOpPattern : public OpRewritePattern<mlir::stablehlo::SqrtOp> {
+  using OpRewritePattern<mlir::stablehlo::SqrtOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(mlir::stablehlo::SqrtOp op,
+                                PatternRewriter& rewriter) const final {
+    TypedAttr lhsAttr;
+    matchPattern(op.getOperand(), m_Constant(&lhsAttr));
+
+    if (!lhsAttr)
+      return rewriter.notifyMatchFailure(op, "operand not constant");
+
+    if (auto res = constFoldUnaryOp<FloatAttr, FloatAttr::ValueType, void>(
+            lhsAttr, foldSqrt)) {
+      rewriter.replaceOpWithNewOp<stablehlo::ConstantOp>(
+          op, op.getType(), llvm::cast<ElementsAttr>(res));
+      return success();
+    }
+
+    return rewriter.notifyMatchFailure(op, "unable to fold sqrt");
+  }
+
+  static std::optional<APFloat> foldSqrt(const APFloat& a) {
+    if (a.getSizeInBits(a.getSemantics()) == 64)
+      return APFloat(std::sqrt(a.convertToDouble()));
+
+    if (a.getSizeInBits(a.getSemantics()) == 32)
+      return APFloat(sqrtf(a.convertToFloat()));
+    return {};
+  }
+};
+
 struct EvalIotaOpPattern : public OpRewritePattern<IotaOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(IotaOp op,
@@ -860,7 +898,7 @@ struct EvalTransposeOpPattern : public OpRewritePattern<TransposeOp> {
   LogicalResult matchAndRewrite(TransposeOp op,
                                 PatternRewriter& rewriter) const override {
     auto resultType = op.getType();
-    if (failed(validateResultTypeForEval(rewriter, op, resultType)))
+    if (failed(validateStaticShapeResult(rewriter, op, resultType)))
       return failure();
 
     ElementsAttr els;
@@ -916,10 +954,9 @@ void populateStablehloAggressiveFolderPatterns(RewritePatternSet* patterns,
 
   // TODO: Consolidate FoldOp patterns
   // One is used by Shape Refinement, the other is a generic folder.
-  patterns
-      ->add<FoldAddOpPattern, FoldBroadcastInDimSplatPattern,
-            FoldConcatenateOpPattern, FoldMulOpPattern, FoldSubtractOpPattern>(
-          context);
+  patterns->add<FoldAddOpPattern, FoldBroadcastInDimSplatPattern,
+                FoldConcatenateOpPattern, FoldMulOpPattern,
+                FoldSubtractOpPattern, FoldSqrtOpPattern>(context);
 }
 
 void populateStablehloShapeFolderPatterns(RewritePatternSet* patterns,
@@ -939,7 +976,7 @@ void populateStablehloShapeFolderPatterns(RewritePatternSet* patterns,
   patterns->add<EvalMulOpPattern>(context, benefit);
   patterns->add<EvalOrOpPattern>(context, benefit);
   patterns->add<EvalRemOpPattern>(context, benefit);
-  patterns->add<EvalReshapeOpPattern>(context, benefit);
+  patterns->add<EvalReshapeOpPattern>(context, benefit, foldFloat);
   patterns->add<EvalSelectOpPattern>(context, benefit);
   patterns->add<EvalSignOpPattern>(context, benefit);
   patterns->add<EvalSliceOpPattern>(context, benefit);

stablehlo/transforms/optimization/StablehloAggressiveSimplificationPatterns.td

@@ -68,6 +68,10 @@ def EmptyI64Array : AttrConstraint<
     CPred<"cast<DenseI64ArrayAttr>($_self).empty()">,
     "is empty i64 array">;
 
+def AnyOne : AttrConstraint<
+    CPred<"::mlir::matchPattern($_self, m_AnyAttrOf(m_One(), m_OneFloat()))">,
+    "is integer one">;
+
 def IntOne : AttrConstraint<
     CPred<"::mlir::matchPattern($_self, m_One())">,
     "is integer one">;
@@ -303,11 +307,11 @@ def : CanonicalizeConstantToRhs<StableHLO_MulOp>;
 
 // Pattern: multiply(X, 0i) -> 0i
 // Multiplication by 0. This fold is not trivial for floats in presence of NaNs
-def : Pat<(StableHLO_MulOp $lhs, (StableHLO_ConstantOp:$zero IntZero:$value)),
+def : Pat<(StableHLO_MulOp $lhs, (StableHLO_ConstantOp:$zero AnyZero:$value)),
           (replaceWithValue $zero)>;
 
 // Pattern: multiply(X, 1i) -> X
-def : Pat<(StableHLO_MulOp $lhs, (StableHLO_ConstantOp IntOne:$value)),
+def : Pat<(StableHLO_MulOp $lhs, (StableHLO_ConstantOp AnyOne:$value)),
           (replaceWithValue $lhs)>;
 
 ////////

stablehlo/transforms/optimization/StablehloTargetIndependentOptimization.cpp

@@ -45,7 +45,7 @@ struct StablehloTargetIndependentOptimizationPass
 
   LogicalResult initialize(MLIRContext* context) override {
     RewritePatternSet patterns_(context);
-    bool foldFloat = false;
+    bool foldFloat = true;
     populateStablehloCanonicalizationPatterns(context, &patterns_);
     populateStablehloAggressiveFolderPatterns(&patterns_, context, foldFloat,
                                               /*benefit=*/2);