[XLA:GPU] Enable chlo.sinh -> kSinh HloInstruction lowering.

ermilovmaxim · tensorflower-gardener · commit adf2e981a312 · 2025-10-06T15:30:04.000-07:00
PiperOrigin-RevId: 815903539
diff --git a/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/optimize_loops.cc b/third_party/xla/xla/backends/gpu/codegen/emitters/transforms/optimize_loops.cc
@@ -54,6 +54,7 @@ bool IsExpensiveToUnroll(mlir::Operation* op) {
       mlir::math::AcosOp,
       mlir::math::AcoshOp,
       mlir::math::AtanhOp,
+      mlir::math::SinhOp,
       mlir::scf::ForOp
       // go/keep-sorted end
       // clang-format on
diff --git a/third_party/xla/xla/hlo/builder/lib/math.cc b/third_party/xla/xla/hlo/builder/lib/math.cc
@@ -1399,7 +1399,11 @@ XlaOp Cosh(XlaOp x, const std::optional<ResultAccuracy>& result_accuracy,
 // +/-89.4159851, due to rounding error when computing x +/- log(1/2).  The
 // correct answer of 3.40281961e+38 (0x7f7fffec) is very close to max-float, so
 // we deem this acceptable.
-XlaOp Sinh(XlaOp x) {
+XlaOp Sinh(XlaOp x, const std::optional<ResultAccuracy>& result_accuracy,
+           bool expand) {
+  if (!expand) {
+    return x.builder()->UnaryOp(HloOpcode::kSinh, x, result_accuracy);
+  }
   XlaBuilder* b = x.builder();
   auto do_it = [&](XlaOp x) -> absl::StatusOr<XlaOp> {
     TF_ASSIGN_OR_RETURN(auto shape, b->GetShape(x));
diff --git a/third_party/xla/xla/hlo/builder/lib/math.h b/third_party/xla/xla/hlo/builder/lib/math.h
@@ -115,7 +115,9 @@ XlaOp Cosh(XlaOp x,
            bool expand = true);
 
 // Computes the hyperbolic sine of 'x'.
-XlaOp Sinh(XlaOp x);
+XlaOp Sinh(XlaOp x,
+           const std::optional<ResultAccuracy>& result_accuracy = std::nullopt,
+           bool expand = true);
 
 // Applies a complex conjugation operation if 'a' is complex and 'conjugate'
 // is true, otherwise returns its argument.
diff --git a/third_party/xla/xla/hlo/builder/xla_builder.h b/third_party/xla/xla/hlo/builder/xla_builder.h
@@ -1772,6 +1772,9 @@ class XlaBuilder {
                     bool expand);
   friend XlaOp Sin(XlaOp operand,
                    const std::optional<ResultAccuracy>& result_accuracy);
+  friend XlaOp Sinh(XlaOp x,
+                    const std::optional<ResultAccuracy>& result_accuracy,
+                    bool expand);
   friend XlaOp Tan(XlaOp operand,
                    const std::optional<ResultAccuracy>& result_accuracy);
   friend XlaOp Tanh(XlaOp operand,
diff --git a/third_party/xla/xla/hlo/ir/hlo_instruction.cc b/third_party/xla/xla/hlo/ir/hlo_instruction.cc
@@ -2759,6 +2759,7 @@ std::unique_ptr<HloInstruction> HloInstruction::CloneWithNewOperands(
     case HloOpcode::kLogistic:
     case HloOpcode::kSign:
     case HloOpcode::kSin:
+    case HloOpcode::kSinh:
     case HloOpcode::kSqrt:
     case HloOpcode::kCbrt:
     case HloOpcode::kTan:
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/gen_hlo_op_writer.td b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/gen_hlo_op_writer.td
@@ -257,6 +257,7 @@ defvar CustomHloConverterOps = [
   MHLO_SendOp,
   MHLO_SetDimensionSizeOp,
   MHLO_SineOp,
+  MHLO_SinhOp,
   MHLO_SortOp,
   MHLO_StochasticConvertOp,
   MHLO_SubtractOp,
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -5246,6 +5246,17 @@ LogicalResult ExportXlaOp(CoshOp op, OpLoweringContext ctx) {
   return success();
 }
 
+LogicalResult ExportXlaOp(SinhOp op, OpLoweringContext ctx) {
+  auto& value_map = *ctx.values;
+  xla::XlaOp operand;
+  if (failed(GetXlaOp(op.getOperand(), value_map, &operand, op))) {
+    return failure();
+  }
+  value_map[op] =
+      xla::Sinh(operand, /*result_accuracy=*/std::nullopt, /*expand=*/false);
+  return success();
+}
+
 LogicalResult ExportXlaOp(AcoshOp op, OpLoweringContext ctx) {
   return ExportElementwiseXlaOp<AcoshOp, xla::Acosh>(op, ctx);
 }
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/BUILD b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/BUILD
@@ -19,6 +19,7 @@ lit_test_suite(
             "case.mlir",
             "composite.mlir",
             "cosh.mlir",
+            "sinh.mlir",
             "dynamic.mlir",
             "export-with-layouts.mlir",
             "export.mlir",
diff --git a/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/sinh.mlir b/third_party/xla/xla/hlo/translate/mhlo_to_hlo/tests/sinh.mlir
@@ -0,0 +1,9 @@
+// RUN: xla-translate -mlir-hlo-to-hlo-text %s | FileCheck %s
+
+module {
+  func.func @main(%arg0: tensor<4xf32>) -> tensor<4xf32> {
+    // CHECK: f32[4] sinh
+    %0 = "mhlo.sinh"(%arg0) : (tensor<4xf32>) -> tensor<4xf32>
+    func.return %0 : tensor<4xf32>
+  }
+}
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc b/third_party/xla/xla/mlir_hlo/mhlo/transforms/chlo_legalize_to_hlo/chlo_legalize_to_hlo_pass.cc
@@ -52,12 +52,14 @@ namespace {
 ChloLegalizeToHighLevelMhloPassOptions FromPassOptions(bool enableAcosh,
                                                        bool enableAcos,
                                                        bool enableAtanh,
-                                                       bool enableCosh) {
+                                                       bool enableCosh,
+                                                       bool enableSinh) {
   ChloLegalizeToHighLevelMhloPassOptions options;
   options.enable_acosh_ = enableAcosh;
   options.enable_acos_ = enableAcos;
   options.enable_atanh_ = enableAtanh;
   options.enable_cosh_ = enableCosh;
+  options.enable_sinh_ = enableSinh;
   return options;
 }
 
@@ -77,6 +79,10 @@ static bool qualifiesForDirectMhloLoweringCosh(chlo::CoshOp op) {
   return llvm::isa<FloatType>(getElementTypeOrSelf(op.getType()));
 }
 
+static bool qualifiesForDirectMhloLoweringSinh(chlo::SinhOp op) {
+  return llvm::isa<FloatType>(getElementTypeOrSelf(op.getType()));
+}
+
 struct ChloLegalizeToHighLevelMhloPass
     : public impl::ChloLegalizeToHighLevelMhloPassBase<
           ChloLegalizeToHighLevelMhloPass> {
@@ -94,7 +100,7 @@ struct ChloLegalizeToHighLevelMhloPass
     chlo::populateChloToHighLevelMhloOpPatterns(
         &context, &conversionPatterns,
         FromPassOptions(enable_acosh_, enable_acos_, enable_atanh_,
-                        enable_cosh_));
+                        enable_cosh_, enable_sinh_));
 
     // Consider the mhlo dialect legal for tests. Also add helper dialects
     // that are needed by the patterns.
@@ -121,6 +127,11 @@ struct ChloLegalizeToHighLevelMhloPass
         return !qualifiesForDirectMhloLoweringCosh(op);
       });
     }
+    if (enable_sinh_) {
+      conversionTarget.addDynamicallyLegalOp<chlo::SinhOp>([](chlo::SinhOp op) {
+        return !qualifiesForDirectMhloLoweringSinh(op);
+      });
+    }
     conversionTarget
         .addIllegalOp<chlo::TopKOp, chlo::ErfOp, chlo::RaggedDotOp>();
 
@@ -254,6 +265,15 @@ LogicalResult convertCoshChloToMhlo(chlo::CoshOp op,
   return success();
 }
 
+LogicalResult convertSinhChloToMhlo(chlo::SinhOp op,
+                                    PatternRewriter& rewriter) {
+  if (!mhlo::qualifiesForDirectMhloLoweringSinh(op)) {
+    return failure();
+  }
+  rewriter.replaceOpWithNewOp<mhlo::SinhOp>(op, op->getOperands());
+  return success();
+}
+
 }  // namespace
 
 ChloLegalizeToHighLevelMhloPassOptions getDefaultChloToHighLevelMhloOptions() {
@@ -266,6 +286,7 @@ ChloLegalizeToHighLevelMhloPassOptions getGpuChloToHighLevelMhloOptions() {
   opts.enable_acos_ = true;
   opts.enable_atanh_ = true;
   opts.enable_cosh_ = true;
+  opts.enable_sinh_ = true;
   return opts;
 }
 
@@ -292,6 +313,9 @@ void populateChloToHighLevelMhloOpPatterns(
   if (options.enable_cosh_) {
     patterns->add(mhlo::convertCoshChloToMhlo, kBenefit);
   }
+  if (options.enable_sinh_) {
+    patterns->add(mhlo::convertSinhChloToMhlo, kBenefit);
+  }
   patterns->add(mhlo::convertRaggedDotChloToMhlo, kBenefit);
   populateWithGenerated(*patterns);
 }
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/hlo_legalize_to_stablehlo/hlo_legalize_to_stablehlo.cc b/third_party/xla/xla/mlir_hlo/mhlo/transforms/hlo_legalize_to_stablehlo/hlo_legalize_to_stablehlo.cc
@@ -130,6 +130,11 @@ std::optional<int64_t> getPublicFeaturesNotInStablehlo(HloOpTy hloOp) {
     // Version 1: Initial version for CoshOp.
     return 1;
   }
+  // StableHLO doesn't support Sinh yet.
+  if constexpr (std::is_same<HloOpTy, mhlo::SinhOp>::value) {
+    // Version 1: Initial version for CoshOp.
+    return 1;
+  }
   return std::nullopt;
 }
 
@@ -463,6 +468,7 @@ LogicalResult convertAttributes(ConversionPatternRewriter& rewriter,
                   !std::is_same<HloOpTy, mhlo::AcoshOp>::value &&
                   !std::is_same<HloOpTy, mhlo::AtanhOp>::value &&
                   !std::is_same<HloOpTy, mhlo::CoshOp>::value &&
+                  !std::is_same<HloOpTy, mhlo::SinhOp>::value &&
                   !std::is_same<HloOpTy, mhlo::ErfOp>::value &&
                   !std::is_same<HloOpTy, mhlo::TopKOp>::value) {
       if (!stablehloAttr) {
@@ -756,10 +762,10 @@ void populateHloToStablehloPatterns(RewritePatternSet* patterns,
       patterns, converter, context, allowExperimentalFeatures,
       allowXlaFeatures);
 
-  populateHloToStablehloCustomCallPatterns<mhlo::AcosOp, mhlo::AcoshOp,
-                                           mhlo::AtanhOp, mhlo::CoshOp,
-                                           mhlo::ErfOp, mhlo::TopKOp>(
-      patterns, converter, context, allowExperimentalFeatures);
+  populateHloToStablehloCustomCallPatterns<
+      mhlo::AcosOp, mhlo::AcoshOp, mhlo::AtanhOp, mhlo::CoshOp, mhlo::SinhOp,
+      mhlo::ErfOp, mhlo::TopKOp>(patterns, converter, context,
+                                 allowExperimentalFeatures);
 }
 
 }  // namespace stablehlo
diff --git a/third_party/xla/xla/mlir_hlo/mhlo/transforms/mhlo_passes.td b/third_party/xla/xla/mlir_hlo/mhlo/transforms/mhlo_passes.td
@@ -30,7 +30,9 @@ def ChloLegalizeToHighLevelMhloPass : Pass<"chlo-legalize-to-high-level-mhlo", "
     Option<"enable_atanh_", "enable-atanh", "bool", /*default=*/"false",
            "Enable chlo.atanh to mhlo.atanh lowering.">,
     Option<"enable_cosh_", "enable-cosh", "bool", /*default=*/"false",
-           "Enable chlo.cosh to mhlo.cosh lowering.">
+           "Enable chlo.cosh to mhlo.cosh lowering.">,
+    Option<"enable_sinh_", "enable-sinh", "bool", /*default=*/"false",
+           "Enable chlo.sinh to mhlo.sinh lowering.">
   ];
   let dependentDialects = ["mhlo::MhloDialect"];
 }
diff --git a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_preserve_high_level_ops.cpp
@@ -266,6 +266,15 @@ struct CoshOpToCustomCallPattern : public OpRewritePattern<chlo::CoshOp> {
   }
 };
 
+struct SinhOpToCustomCallPattern : public OpRewritePattern<chlo::SinhOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(chlo::SinhOp op,
+                                PatternRewriter& rewriter) const override {
+    return wrapChloOperationInCustomCall(rewriter, op, "mhlo.sinh",
+                                         /*version=*/1);
+  }
+};
+
 ///////
 // CHLO to CompositeOp Patterns
 ///////
@@ -341,6 +350,14 @@ struct CoshOpToCompositePattern : public OpRewritePattern<chlo::CoshOp> {
   }
 };
 
+struct SinhOpToCompositePattern : public OpRewritePattern<chlo::SinhOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(chlo::SinhOp op,
+                                PatternRewriter& rewriter) const override {
+    return wrapChloOpInComposite(op, /*version=*/1, rewriter);
+  }
+};
+
 }  // namespace
 
 struct ChloPreserveHighLevelOpsPass
@@ -368,6 +385,7 @@ struct ChloPreserveHighLevelOpsPass
         AcoshOpToCustomCallPattern,
         AtanhOpToCustomCallPattern,
         CoshOpToCustomCallPattern,
+        SinhOpToCustomCallPattern,
         ErfOpToCustomCallPattern,
         RaggedDotOpToCustomCallPattern,
         TopKOpToCustomCallPattern>(ctx);
@@ -377,6 +395,7 @@ struct ChloPreserveHighLevelOpsPass
         AcoshOpToCompositePattern,
         AtanhOpToCompositePattern,
         CoshOpToCompositePattern,
+        SinhOpToCompositePattern,
         ErfOpToCompositePattern,
         RaggedDotOpToCompositePattern,
         TopKOpToCompositePattern>(ctx);
diff --git a/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp b/third_party/xla/xla/mlir_hlo/stablehlo_ext/transforms/chlo_recompose_ops.cpp
@@ -312,6 +312,22 @@ struct CoshOpRecomposePattern
   }
 };
 
+struct SinhOpRecomposePattern
+    : public OpRewritePattern<stablehlo::CompositeOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(stablehlo::CompositeOp op,
+                                PatternRewriter& rewriter) const override {
+    if (op.getName() != "chlo.sinh") {
+      return rewriter.notifyMatchFailure(op, "not a chlo.sinh");
+    }
+    if (op.getVersion() != 1) {
+      return rewriter.notifyMatchFailure(
+          op, "unsupported version for chlo.cosh composite");
+    }
+    return recomposeChloOpFromCompositeOp<chlo::SinhOp>(op, rewriter);
+  }
+};
+
 struct ErfOpRecomposePattern : public OpRewritePattern<stablehlo::CompositeOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(stablehlo::CompositeOp op,
@@ -441,6 +457,16 @@ struct CoshOpCustomCallRecomposePattern
   }
 };
 
+struct SinhOpCustomCallRecomposePattern
+    : public OpRewritePattern<stablehlo::CustomCallOp> {
+  using OpRewritePattern::OpRewritePattern;
+  LogicalResult matchAndRewrite(stablehlo::CustomCallOp op,
+                                PatternRewriter& rewriter) const override {
+    return recomposeChloOpFromCustomCall<chlo::SinhOp>(
+        op, {"mhlo.sinh", "chlo.sinh"}, rewriter);
+  }
+};
+
 }  // namespace
 
 struct ChloRecomposeOpsPass
@@ -465,6 +491,7 @@ struct ChloRecomposeOpsPass
       AcoshOpCustomCallRecomposePattern,
       AtanhOpCustomCallRecomposePattern,
       CoshOpCustomCallRecomposePattern,
+      SinhOpCustomCallRecomposePattern,
       ErfOpCustomCallRecomposePattern,
       RaggedDotOpCustomCallRecomposePattern,
       TanOpCustomCallRecomposePattern,
@@ -476,6 +503,7 @@ struct ChloRecomposeOpsPass
       AcoshOpRecomposePattern,
       AtanhOpRecomposePattern,
       CoshOpRecomposePattern,
+      SinhOpRecomposePattern,
       ErfOpRecomposePattern,
       RaggedDotOpRecomposePattern,
       TopKOpRecomposePattern>(ctx);
diff --git a/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir b/third_party/xla/xla/mlir_hlo/tests/Dialect/chlo/chlo_legalize_to_mhlo.mlir
@@ -1,5 +1,5 @@
 // RUN: mlir-hlo-opt --chlo-legalize-to-hlo --split-input-file -verify-diagnostics %s | FileCheck %s --dump-input-context=20
-// RUN: mlir-hlo-opt --chlo-legalize-to-high-level-mhlo="enable-acosh enable-acos enable-atanh enable-cosh" --split-input-file -verify-diagnostics %s | FileCheck %s --check-prefix=CHECK-HIGH-LEVEL
+// RUN: mlir-hlo-opt --chlo-legalize-to-high-level-mhlo="enable-acosh enable-acos enable-atanh enable-cosh enable-sinh" --split-input-file -verify-diagnostics %s | FileCheck %s --check-prefix=CHECK-HIGH-LEVEL
 
 // CHECK-LABEL: func.func @asin_bf16(
 // CHECK-SAME:    %[[TMP_arg0:.*]]: tensor<bf16>
@@ -2640,6 +2640,7 @@ func.func @sinh_f32(%x : tensor<f32>) -> tensor<f32> {
   %1 = chlo.sinh %x : tensor<f32> -> tensor<f32>
   func.return %1 : tensor<f32>
 }
+// CHECK-HIGH-LEVEL: mhlo.sinh
 
 // -----
 
@@ -2652,6 +2653,7 @@ func.func @sinh_f16(%x : tensor<f16>) -> tensor<f16> {
   %1 = chlo.sinh %x : tensor<f16> -> tensor<f16>
   func.return %1 : tensor<f16>
 }
+// CHECK-HIGH-LEVEL: mhlo.sinh
 
 // -----
 
@@ -2669,6 +2671,7 @@ func.func @sinh_complex(%x : tensor<2xcomplex<f32>>) -> tensor<2xcomplex<f32>> {
   %1 = chlo.sinh %x : tensor<2xcomplex<f32>> -> tensor<2xcomplex<f32>>
   func.return %1 : tensor<2xcomplex<f32>>
 }
+// CHECK-HIGH-LEVEL-NOT: mhlo.sinh
 
 // -----
 
diff --git a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_preserve_high_level_ops.mlir
@@ -113,6 +113,16 @@ func.func @cosh_preserve(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
 
 // -----
 
+// CHECK-LABEL: func @sinh_preserve
+func.func @sinh_preserve(%arg0: tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16> {
+  // CHECK-CC: stablehlo.custom_call @mhlo.sinh(%arg0) {mhlo.attributes = {}, mhlo.version = 1 : i64} : (tensor<3x20x20xbf16>) -> tensor<?x20x20xbf16>
+  // CHECK: stablehlo.composite "chlo.sinh" %arg0 {decomposition = @chlo.sinh.impl, version = 1 : i32}
+  %0 = chlo.sinh %arg0 : tensor<3x20x20xbf16> -> tensor<?x20x20xbf16>
+  return %0 : tensor<?x20x20xbf16>
+}
+
+// -----
+
 // CHECK-LABEL: func @tan_no_preserve
 func.func @tan_no_preserve(%arg0: tensor<16xf32>) -> tensor<?xf32> {
   // CHECK: chlo.tan
diff --git a/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir b/third_party/xla/xla/mlir_hlo/tests/stablehlo_ext/chlo_recompose_ops.mlir
diff --git a/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc b/third_party/xla/xla/tests/exhaustive/exhaustive_unary_test_ops.inc
