Add StableHLO complex log to stablehlo-complex-math-expander pass

build_tools/math/README.md

@@ -31,7 +31,7 @@ following requirements:
 
 - Python 3.11 or newer
 - mpmath 1.3 or newer
-- functional_algorithms 0.12 or newer
+- functional_algorithms 0.14.1 or newer
 
 that can be installed via pypi:
 

build_tools/math/generate_ChloDecompositionPatternsMath.py

@@ -98,10 +98,17 @@ def main(kind="CHLO"):
       ("CHLO_SquareOp", "complex_square", ("z:complex",)),
       ("CHLO_SquareOp", "real_square", ("x:float",)),
       ("StableHLO_Log1pOp", "complex_log1p", ("z:complex",)),
+      ("StableHLO_SqrtOp", "complex_sqrt", ("z:complex",)),
+      ("StableHLO_LogOp", "complex_log", ("z:complex",)),
   ]:
     if not chloname.startswith(kind):
       continue
-    print(f'Generating {chloname} from {fname}{args}')
+    if chloname.startswith("StableHLO_"):
+      NameOp = chloname.split("_", 1)[1]
+      expander_name = f"{NameOp}_ComplexElementType_ComplexMathExpander"
+    else:
+      expander_name = ""
+    print(f"Generating {chloname} from {fname}{args}")
     func = getattr(fa.algorithms, fname, None)
     if func is None:
       warnings.warn(
@@ -110,22 +117,12 @@ def main(kind="CHLO"):
     ctx = fa.Context(paths=[fa.algorithms],
                      parameters=dict(rewrite_keep_integer_literals=True))
     graph = ctx.trace(func, *args).rewrite(target, fa.rewrite)
-    graph.props.update(name=chloname)
+    graph.props.update(name=chloname, expander_name=expander_name)
     src = graph.tostring(target)
     sources.append(target.make_comment(func.__doc__)) if func.__doc__ else None
     sources[-1] += src
   source = "\n\n".join(sources) + "\n"
 
-  if chloname.startswith('StableHLO_'):
-    # an ugly hack to fix the definition of stablehlo complex math
-    # functions. TODO(pearu): add the corresponding feature to
-    # functional_algorithms stablehlo printer
-    NameOp = chloname.split('_', 1)[1]
-    source = source.replace(
-        f'def : Pat<({chloname}',
-        f'def {NameOp}_ComplexElementType_ComplexMathExpander : Pat<({chloname}'
-    )
-
   if os.path.isfile(output_file):
     f = open(output_file, "r")
     content = f.read()
@@ -177,6 +174,8 @@ def ComplexElementType : Type<
 def StableHLO_ConstantLikeMaxFiniteValue : NativeCodeCall<
     "::mlir::stablehlo::getConstantLikeMaxFiniteValue($_builder, $_loc, $0)">;
 
+def StableHLO_ConstantLikePosInfValue : NativeCodeCall<
+    "::mlir::stablehlo::getConstantLikeInfValue($_builder, $_loc, $0, /*negative=*/false)">;
 """)
   f.write(source)
   f.close()

build_tools/math/generate_tests.py

@@ -54,6 +54,10 @@
     #   that defines the precision of computing reference values:
     #   mpmath.mp.prec * extra_prec_multiplier
     #
+    #   namespace - "chlo" or "stablehlo"
+    #
+    #   passes - a string of pass arguments
+    #
     # When unspecifed, these parameters are retrieved from
     # functional_algorithms database of support functions.
     #
@@ -68,6 +72,14 @@
          mpmath_name="log1p",
          namespace="stablehlo",
          passes="--stablehlo-complex-math-expander"),
+    dict(name="log",
+         mpmath_name="log",
+         namespace="stablehlo",
+         passes="--stablehlo-complex-math-expander"),
+    dict(name="sqrt",
+         mpmath_name="sqrt",
+         namespace="stablehlo",
+         passes="--stablehlo-complex-math-expander"),
 ]
 
 
@@ -109,8 +121,8 @@ def main():
 
   if fa_version < required_fa_version:
     msg = (
-        f"functional_algorithm version {'.'.join(map(str, required_fa_version))}"
-        f" or newer is required, got {fa.__version__}")
+        f"functional_algorithms version {'.'.join(map(str, required_fa_version))}"
+        f" or newer is required, got {fa.__version__}, stopping.")
     warnings.warn(msg)
     return
 

stablehlo/tests/math/log_complex128.mlir

@@ -0,0 +1,19 @@
+// RUN: stablehlo-opt --stablehlo-complex-math-expander %s | stablehlo-translate --interpret
+// This file is generated, see build_tools/math/README.md for more information.
+module @log_complex128 {
+  func.func private @samples() -> tensor<169xcomplex<f64>> {
+    %0 = stablehlo.constant dense<"0x000000000000F0FF000000000000F0FFFFFFFFFFFFFFEFFF000000000000F0FFFEFFFFFFFFFFEFFF000000000000F0FF000000000000F8BF000000000000F0FF000000000000FC9F000000000000F0FF0100000000000080000000000000F0FF0000000000000000000000000000F0FF0100000000000000000000000000F0FF000000000000FC1F000000000000F0FF000000000000F83F000000000000F0FFFEFFFFFFFFFFEF7F000000000000F0FFFFFFFFFFFFFFEF7F000000000000F0FF000000000000F07F000000000000F0FF000000000000F0FFFFFFFFFFFFFFEFFFFFFFFFFFFFFFEFFFFFFFFFFFFFFFEFFFFEFFFFFFFFFFEFFFFFFFFFFFFFFFEFFF000000000000F8BFFFFFFFFFFFFFEFFF000000000000FC9FFFFFFFFFFFFFEFFF0100000000000080FFFFFFFFFFFFEFFF0000000000000000FFFFFFFFFFFFEFFF0100000000000000FFFFFFFFFFFFEFFF000000000000FC1FFFFFFFFFFFFFEFFF000000000000F83FFFFFFFFFFFFFEFFFFEFFFFFFFFFFEF7FFFFFFFFFFFFFEFFFFFFFFFFFFFFFEF7FFFFFFFFFFFFFEFFF000000000000F07FFFFFFFFFFFFFEFFF000000000000F0FFFEFFFFFFFFFFEFFFFFFFFFFFFFFFEFFFFEFFFFFFFFFFEFFFFEFFFFFFFFFFEFFFFEFFFFFFFFFFEFFF000000000000F8BFFEFFFFFFFFFFEFFF000000000000FC9FFEFFFFFFFFFFEFFF0100000000000080FEFFFFFFFFFFEFFF0000000000000000FEFFFFFFFFFFEFFF0100000000000000FEFFFFFFFFFFEFFF000000000000FC1FFEFFFFFFFFFFEFFF000000000000F83FFEFFFFFFFFFFEFFFFEFFFFFFFFFFEF7FFEFFFFFFFFFFEFFFFFFFFFFFFFFFEF7FFEFFFFFFFFFFEFFF000000000000F07FFEFFFFFFFFFFEFFF000000000000F0FF000000000000F8BFFFFFFFFFFFFFEFFF000000000000F8BFFEFFFFFFFFFFEFFF000000000000F8BF000000000000F8BF000000000000F8BF000000000000FC9F000000000000F8BF0100000000000080000000000000F8BF0000000000000000000000000000F8BF0100000000000000000000000000F8BF000000000000FC1F000000000000F8BF000000000000F83F000000000000F8BFFEFFFFFFFFFFEF7F000000000000F8BFFFFFFFFFFFFFEF7F000000000000F8BF000000000000F07F000000000000F8BF000000000000F0FF000000000000FC9FFFFFFFFFFFFFEFFF000000000000FC9FFEFFFFFFFFFFEFFF000000000000FC9F000000000000F8BF000000000000FC9F000000000000FC9F000000000000FC9F0100000000000080000000000000FC9F0000000000000000000000000000FC9F0100000000000000000000000000FC9F000000000000FC1F000000000000FC9F000000000000F83F000000000000FC9FFEFFFFFFFFFFEF7F000000000000FC9FFFFFFFFFFFFFEF7F000000000000FC9F000000000000F07F000000000000FC9F000000000000F0FF0100000000000080FFFFFFFFFFFFEFFF0100000000000080FEFFFFFFFFFFEFFF0100000000000080000000000000F8BF0100000000000080000000000000FC9F0100000000000080010000000000008001000000000000800000000000000000010000000000008001000000000000000100000000000080000000000000FC1F0100000000000080000000000000F83F0100000000000080FEFFFFFFFFFFEF7F0100000000000080FFFFFFFFFFFFEF7F0100000000000080000000000000F07F0100000000000080000000000000F0FF0000000000000000FFFFFFFFFFFFEFFF0000000000000000FEFFFFFFFFFFEFFF0000000000000000000000000000F8BF0000000000000000000000000000FC9F0000000000000000010000000000008000000000000000000000000000000000000000000000000001000000000000000000000000000000000000000000FC1F0000000000000000000000000000F83F0000000000000000FEFFFFFFFFFFEF7F0000000000000000FFFFFFFFFFFFEF7F0000000000000000000000000000F07F0000000000000000000000000000F0FF0100000000000000FFFFFFFFFFFFEFFF0100000000000000FEFFFFFFFFFFEFFF0100000000000000000000000000F8BF0100000000000000000000000000FC9F0100000000000000010000000000008001000000000000000000000000000000010000000000000001000000000000000100000000000000000000000000FC1F0100000000000000000000000000F83F0100000000000000FEFFFFFFFFFFEF7F0100000000000000FFFFFFFFFFFFEF7F0100000000000000000000000000F07F0100000000000000000000000000F0FF000000000000FC1FFFFFFFFFFFFFEFFF000000000000FC1FFEFFFFFFFFFFEFFF000000000000FC1F000000000000F8BF000000000000FC1F000000000000FC9F000000000000FC1F0100000000000080000000000000FC1F0000000000000000000000000000FC1F0100000000000000000000000000FC1F000000000000FC1F000000000000FC1F000000000000F83F000000000000FC1FFEFFFFFFFFFFEF7F000000000000FC1FFFFFFFFFFFFFEF7F000000000000FC1F000000000000F07F000000000000FC1F000000000000F0FF000000000000F83FFFFFFFFFFFFFEFFF000000000000F83FFEFFFFFFFFFFEFFF000000000000F83F000000000000F8BF000000000000F83F000000000000FC9F000000000000F83F0100000000000080000000000000F83F0000000000000000000000000000F83F0100000000000000000000000000F83F000000000000FC1F000000000000F83F000000000000F83F000000000000F83FFEFFFFFFFFFFEF7F000000000000F83FFFFFFFFFFFFFEF7F000000000000F83F000000000000F07F000000000000F83F000000000000F0FFFEFFFFFFFFFFEF7FFFFFFFFFFFFFEFFFFEFFFFFFFFFFEF7FFEFFFFFFFFFFEFFFFEFFFFFFFFFFEF7F000000000000F8BFFEFFFFFFFFFFEF7F000000000000FC9FFEFFFFFFFFFFEF7F0100000000000080FEFFFFFFFFFFEF7F0000000000000000FEFFFFFFFFFFEF7F0100000000000000FEFFFFFFFFFFEF7F000000000000FC1FFEFFFFFFFFFFEF7F000000000000F83FFEFFFFFFFFFFEF7FFEFFFFFFFFFFEF7FFEFFFFFFFFFFEF7FFFFFFFFFFFFFEF7FFEFFFFFFFFFFEF7F000000000000F07FFEFFFFFFFFFFEF7F000000000000F0FFFFFFFFFFFFFFEF7FFFFFFFFFFFFFEFFFFFFFFFFFFFFFEF7FFEFFFFFFFFFFEFFFFFFFFFFFFFFFEF7F000000000000F8BFFFFFFFFFFFFFEF7F000000000000FC9FFFFFFFFFFFFFEF7F0100000000000080FFFFFFFFFFFFEF7F0000000000000000FFFFFFFFFFFFEF7F0100000000000000FFFFFFFFFFFFEF7F000000000000FC1FFFFFFFFFFFFFEF7F000000000000F83FFFFFFFFFFFFFEF7FFEFFFFFFFFFFEF7FFFFFFFFFFFFFEF7FFFFFFFFFFFFFEF7FFFFFFFFFFFFFEF7F000000000000F07FFFFFFFFFFFFFEF7F000000000000F0FF000000000000F07FFFFFFFFFFFFFEFFF000000000000F07FFEFFFFFFFFFFEFFF000000000000F07F000000000000F8BF000000000000F07F000000000000FC9F000000000000F07F0100000000000080000000000000F07F0000000000000000000000000000F07F0100000000000000000000000000F07F000000000000FC1F000000000000F07F000000000000F83F000000000000F07FFEFFFFFFFFFFEF7F000000000000F07FFFFFFFFFFFFFEF7F000000000000F07F000000000000F07F000000000000F07F"> : tensor<169xcomplex<f64>>
+    return %0 : tensor<169xcomplex<f64>>
+  }
+  func.func private @expected() -> tensor<169xcomplex<f64>> {
+    %0 = stablehlo.constant dense<"0x000000000000F07FD221337F7CD902C0000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21F9BF000000000000F07F182D4454FB21E9BF000000000000F07F182D4454FB2109C036195AC708318640D221337F7CD902C036195AC708318640D221337F7CD902C0EF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BF36195AC708318640192D4454FB21E9BF36195AC708318640182D4454FB21E9BF000000000000F07F0000000000000000000000000000F07F182D4454FB2109C036195AC708318640D221337F7CD902C036195AC708318640D221337F7CD902C0EF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BFEF39FAFE422E8640182D4454FB21F9BF36195AC708318640182D4454FB21E9BF36195AC708318640182D4454FB21E9BF000000000000F07F0000000000000000000000000000F07F182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C01EE9DC75B310E83FD221337F7CD902C04C98BFEC23F3D93F182D4454FB21F9BF4C98BFEC23F3D93F182D4454FB21F9BF4C98BFEC23F3D93F182D4454FB21F9BF4C98BFEC23F3D93F182D4454FB21F9BF4C98BFEC23F3D93F182D4454FB21F9BF1EE9DC75B310E83F182D4454FB21E9BFEF39FAFE422E86400000000000000680EF39FAFE422E86400000000000000680000000000000F07F0000000000000000000000000000F07F182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C04C98BFEC23F3D93F182D4454FB2109C040438B3EC31F76C0D221337F7CD902C0CF014BCF4E2576C0182D4454FB21F9BFCF014BCF4E2576C0182D4454FB21F9BFCF014BCF4E2576C0182D4454FB21F9BF40438B3EC31F76C0182D4454FB21E9BF4C98BFEC23F3D93FABAAAAAAAAAAF29FEF39FAFE422E86400000000000000080EF39FAFE422E86400000000000000080000000000000F07F0000000000000000000000000000F07F182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C0EF39FAFE422E8640182D4454FB2109C04C98BFEC23F3D93F182D4454FB2109C0CF014BCF4E2576C0182D4454FB2109C07C920D7CBF4087C0D221337F7CD902C0C3716D44854387C0182D4454FB21F9BF7C920D7CBF4087C0182D4454FB21E9BFCF014BCF4E2576C0922449922449C29C4C98BFEC23F3D93F0000000000000080EF39FAFE422E86400000000000000080EF39FAFE422E86400000000000000080000000000000F07F0000000000000000000000000000F07F182D4454FB210940EF39FAFE422E8640182D4454FB210940EF39FAFE422E8640182D4454FB2109404C98BFEC23F3D93F182D4454FB210940CF014BCF4E2576C0182D4454FB210940C3716D44854387C0182D4454FB210940000000000000F0FF0000000000000000C3716D44854387C00000000000000000CF014BCF4E2576C000000000000000004C98BFEC23F3D93F0000000000000000EF39FAFE422E86400000000000000000EF39FAFE422E86400000000000000000000000000000F07F0000000000000000000000000000F07F182D4454FB210940EF39FAFE422E8640182D4454FB210940EF39FAFE422E8640182D4454FB2109404C98BFEC23F3D93F182D4454FB210940CF014BCF4E2576C0182D4454FB2109407C920D7CBF4087C0D221337F7CD90240C3716D44854387C0182D4454FB21F93F7C920D7CBF4087C0182D4454FB21E93FCF014BCF4E2576C0922449922449C21C4C98BFEC23F3D93F0000000000000000EF39FAFE422E86400000000000000000EF39FAFE422E86400000000000000000000000000000F07F0000000000000000000000000000F07F182D4454FB210940EF39FAFE422E8640182D4454FB210940EF39FAFE422E8640182D4454FB2109404C98BFEC23F3D93F182D4454FB21094040438B3EC31F76C0D221337F7CD90240CF014BCF4E2576C0182D4454FB21F93FCF014BCF4E2576C0182D4454FB21F93FCF014BCF4E2576C0182D4454FB21F93F40438B3EC31F76C0182D4454FB21E93F4C98BFEC23F3D93FABAAAAAAAAAAF21FEF39FAFE422E86400000000000000000EF39FAFE422E86400000000000000000000000000000F07F0000000000000000000000000000F07F182D4454FB210940EF39FAFE422E8640182D4454FB210940EF39FAFE422E8640182D4454FB2109401EE9DC75B310E83FD221337F7CD902404C98BFEC23F3D93F182D4454FB21F93F4C98BFEC23F3D93F182D4454FB21F93F4C98BFEC23F3D93F182D4454FB21F93F4C98BFEC23F3D93F182D4454FB21F93F4C98BFEC23F3D93F182D4454FB21F93F1EE9DC75B310E83F182D4454FB21E93FEF39FAFE422E86400000000000000600EF39FAFE422E86400000000000000600000000000000F07F0000000000000000000000000000F07F182D4454FB21094036195AC708318640D221337F7CD9024036195AC708318640D221337F7CD90240EF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93F36195AC708318640182D4454FB21E93F36195AC708318640182D4454FB21E93F000000000000F07F0000000000000000000000000000F07F182D4454FB21094036195AC708318640D221337F7CD9024036195AC708318640D221337F7CD90240EF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93FEF39FAFE422E8640182D4454FB21F93F36195AC708318640192D4454FB21E93F36195AC708318640182D4454FB21E93F000000000000F07F0000000000000000000000000000F07FD221337F7CD90240000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21F93F000000000000F07F182D4454FB21E93F"> : tensor<169xcomplex<f64>>
+    return %0 : tensor<169xcomplex<f64>>
+  }
+  func.func public @main() {
+    %0 = call @samples() : () -> tensor<169xcomplex<f64>>
+    %1 = "stablehlo.log"(%0) : (tensor<169xcomplex<f64>>) -> tensor<169xcomplex<f64>>
+    %2 = call @expected() : () -> tensor<169xcomplex<f64>>
+    check.expect_close %1, %2, max_ulp_difference = 3 : tensor<169xcomplex<f64>>, tensor<169xcomplex<f64>>
+    func.return
+  }
+}