Add stablehlo-complex-math-expander pass.

BUILD.bazel

@@ -370,6 +370,21 @@ gentbl_cc_library(
     ],
 )
 
+gentbl_cc_library(
+    name = "stablehlo_create_complex_math_expander_inc_gen",
+    tbl_outs = [
+        (
+            ["--gen-rewriters"],
+            "stablehlo/transforms/StablehloComplexMathExpanderPatterns.h.inc",
+        ),
+    ],
+    tblgen = "@llvm-project//mlir:mlir-tblgen",
+    td_file = "stablehlo/transforms/StablehloComplexMathExpanderPatterns.td",
+    deps = [
+        ":stablehlo_ops_td_files",
+    ],
+)
+
 cc_library(
     name = "interpreter_ops",
     srcs = [
@@ -1120,6 +1135,7 @@ cc_library(
         "stablehlo/transforms/StablehloAggressiveSimplification.cpp",
         "stablehlo/transforms/StablehloCanonicalizeDynamism.cpp",
         "stablehlo/transforms/StablehloCompatibilityExpander.cpp",
+        "stablehlo/transforms/StablehloComplexMathExpander.cpp",
         "stablehlo/transforms/StablehloConvertToSignless.cpp",
         "stablehlo/transforms/StablehloLegalizeCompositeToCall.cpp",
         "stablehlo/transforms/StablehloLegalizeDeprecatedOps.cpp",
@@ -1148,6 +1164,7 @@ cc_library(
         ":linalg_passes",
         ":stablehlo_aggressive_simplification_inc_gen",
         ":stablehlo_create_compatibility_expander_inc_gen",
+        ":stablehlo_create_complex_math_expander_inc_gen",
         ":stablehlo_legalize_deprecated_ops_inc_gen",
         ":stablehlo_ops",
         ":stablehlo_ops_inc_gen",

stablehlo/transforms/CMakeLists.txt

@@ -28,6 +28,10 @@ set(LLVM_TARGET_DEFINITIONS StablehloCompatibilityExpanderPatterns.td)
 mlir_tablegen(StablehloCompatibilityExpanderPatterns.h.inc --gen-rewriters)
 add_public_tablegen_target(StablehloCompatibilityExpanderPatternsIncGen)
 
+set(LLVM_TARGET_DEFINITIONS StablehloComplexMathExpanderPatterns.td)
+mlir_tablegen(StablehloComplexMathExpanderPatterns.h.inc --gen-rewriters)
+add_public_tablegen_target(StablehloComplexMathExpanderPatternsIncGen)
+
 set(LLVM_TARGET_DEFINITIONS StablehloLegalizeDeprecatedOpsPatterns.td)
 mlir_tablegen(StablehloLegalizeDeprecatedOpsPatterns.h.inc --gen-rewriters)
 add_public_tablegen_target(StablehloLegalizeDeprecatedOpsPatternsIncGen)
@@ -47,6 +51,7 @@ add_mlir_dialect_library(StablehloPasses
   StablehloCanonicalizeDynamism.cpp
   StablehloConvertToSignless.cpp
   StablehloCompatibilityExpander.cpp
+  StablehloComplexMathExpander.cpp
   StablehloLegalizeCompositeToCall.cpp
   StablehloLegalizeDeprecatedOps.cpp
   StablehloLegalizeQuantToMath.cpp
@@ -64,6 +69,7 @@ add_mlir_dialect_library(StablehloPasses
   PassesIncGen
   StablehloAggressiveSimplificationPatternsIncGen
   StablehloCompatibilityExpanderPatternsIncGen
+  StablehloComplexMathExpanderPatternsIncGen
   StablehloLegalizeDeprecatedOpsPatternsIncGen
   VhloToVersionPatterns
 

stablehlo/transforms/Passes.h

@@ -144,6 +144,12 @@ void createStablehloRemoveDynamismPipeline(OpPassManager &pm,
 // operations into a primitive math operations.
 void createStablehloLowerQuantPipeline(OpPassManager &pm);
 
+/// Collection of patterns to create expander for StableHLO complex
+/// math operations.
+void populateStablehloComplexMathExpanderPatterns(
+    RewritePatternSet *patterns, MLIRContext *context/*,
+                                                       vhlo::Version targetVersion*/);
+
 // Adds `stablehlo-deserialize` pipeline as a registered pass pipeline
 // for opt tools.
 void registerPassPipelines();

stablehlo/transforms/Passes.td

@@ -123,6 +123,40 @@ def StablehloCompatibilityExpanderPass : Pass<"stablehlo-compatibility-expander"
   ];
 }
 
+def StablehloComplexMathExpanderPass : Pass<"stablehlo-complex-math-expander", "mlir::func::FuncOp"> {
+  let summary = "Expander for StableHLO complex math operations.";
+
+  let description = [{
+    StableHLO complex math operations are decompositions using
+    StableHLO real math operations.
+
+    This statement is based on the assumption that no hardware exists
+    that supports complex numbers nor complex math operations
+    natively. This means that the fallback mechanisms on complex math
+    operations that compilers may implement, are redundant. With
+    enabling this pass, all StableHLO complex math operations will be
+    expanded.
+
+    ```mlir
+    func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
+      %1 = stablehlo.sqrt %arg0 : tensor<4xcomplex<f64>>
+      func.return %1 : tensor<4xcomplex<f64>>
+    }
+
+    ==>
+
+    func.func @sqrt_op_complex(%arg0: tensor<4xcomplex<f64>>) -> tensor<4xcomplex<f64>> {
+      TBD
+      return %2 : tensor<4xcomplex<f64>>
+    }
+    ```
+  }];
+  let dependentDialects = [
+    "mlir::stablehlo::StablehloDialect",
+    "mlir::chlo::ChloDialect",
+  ];
+}
+
 def StablehloConvertToSignlessPass : Pass<"stablehlo-convert-to-signless", "ModuleOp"> {
   let summary = "Pass to transform the IR to be on signless integers.";
 }

stablehlo/transforms/StablehloComplexMathExpander.cpp

@@ -0,0 +1,81 @@
+/* Copyright 2024 The StableHLO Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include <memory>
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Pass/Pass.h"
+#include "stablehlo/dialect/ChloOps.h"
+#include "stablehlo/dialect/StablehloOps.h"
+
+#include "llvm/ADT/APFloat.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypeInterfaces.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Rewrite/FrozenRewritePatternSet.h"
+#include "mlir/Support/LLVM.h"
+#include "mlir/Transforms/DialectConversion.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "stablehlo/dialect/StablehloOps.h"
+#include "stablehlo/dialect/Version.h"
+#include "stablehlo/transforms/PassUtils.h"
+#include "stablehlo/transforms/Passes.h"
+
+namespace mlir {
+namespace stablehlo {
+#define GEN_PASS_DEF_STABLEHLOCOMPLEXMATHEXPANDERPASS
+#include "stablehlo/transforms/Passes.h.inc"
+
+namespace {
+
+static Value getConstantLikeInfValue(OpBuilder &b, Location loc, Value val,
+                                     bool negative) {
+  auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
+  return getConstantLike(
+      b, loc, llvm::APFloat::getInf(ty.getFloatSemantics(), negative), val);
+}
+
+static Value getConstantLikeMaxFiniteValue(OpBuilder &b, Location loc,
+                                           Value val) {
+  auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
+  return getConstantLike(
+      b, loc, llvm::APFloat::getLargest(ty.getFloatSemantics()), val);
+}
+
+static Value getConstantLikeSmallestNormalizedValue(OpBuilder &b, Location loc,
+                                                    Value val) {
+  auto ty = cast<FloatType>(getElementTypeOrSelf(val.getType()));
+  return getConstantLike(
+      b, loc, llvm::APFloat::getSmallestNormalized(ty.getFloatSemantics()),
+      val);
+}
+
+#include "stablehlo/transforms/StablehloComplexMathExpanderPatterns.h.inc"
+
+}  // namespace
+
+void populateStablehloComplexMathExpanderPatterns(
+    RewritePatternSet *patterns, MLIRContext *context/*,
+                                                       vhlo::Version targetVersion*/) {
+  // StableHLO Log1pOp is introduced in v0.9.0.
+  patterns->add<Log1pOp_ComplexElementType_ComplexMathExpander>(context);
+}
+
+}  // namespace stablehlo
+}  // namespace mlir

stablehlo/transforms/StablehloComplexMathExpanderPatterns.td

@@ -0,0 +1,18 @@
+/* Copyright 2024 The StableHLO Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+include "ChloDecompositionPatterns.td"
+
+def Log1pOp_ComplexElementType_ComplexMathExpander : Pat<(StableHLO_Log1pOp ComplexElementType:$input), (CHLO_Log1pOp $input)>;