// Copyright 2020 The IREE Authors
//
// Licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

// This is the legalization pattern definition file for CHLO to StableHLO.
// These are included in the populateDecompositionPatterns factory
// and should only include canonical expansions which are not actually
// ambiguous/different for various backends. Avoid patterns that are actually
// lowering to non-canonical forms.

include "mlir/IR/OpBase.td"
include "stablehlo/dialect/ChloOps.td"
include "stablehlo/dialect/StablehloOps.td"

class StableHLO_ComparisonDirectionValue<string enumStr> :
  ConstantAttr<StableHLO_ComparisonDirectionAttr,
               "::mlir::stablehlo::ComparisonDirection::" # enumStr>;

class ConstantLike<string value> : NativeCodeCall<
    "::mlir::iree_compiler::stablehlo::getConstantLike($_builder, $_loc, " # value # ", $0)">;

def ComplexElementType : Type<
  CPred<"isa<ComplexType>(cast<ShapedType>($_self).getElementType())">,
  "Complex element type">;

def NonComplexElementType : Type<
  CPred<"!isa<ComplexType>(cast<ShapedType>($_self).getElementType())">,
  "Non-complex element type">;

def ConstantLikeMaxFiniteValue : NativeCodeCall<
    "::mlir::iree_compiler::stablehlo::getConstantLikeMaxFiniteValue($_builder, $_loc, $0)">;

def ConstantLikePosInfValue : NativeCodeCall<
    "::mlir::iree_compiler::stablehlo::getConstantLikeInfValue($_builder, $_loc, $0, /*negative=*/false)">;

def ConstantLikeNegInfValue : NativeCodeCall<
    "::mlir::iree_compiler::stablehlo::getConstantLikeInfValue($_builder, $_loc, $0, /*negative=*/true)">;

//===----------------------------------------------------------------------===//
// Unary op patterns.
//===----------------------------------------------------------------------===//

// Expand acos for non-complex arguments to MHLO dialect as follows:
//   acos(x) = 2 * atan2(sqrt(1 - x^2), (1 + x))  if x != -1
//           = pi                                 if x == -1
//
// TODO(b/237376133): Support operands with complex element types separately
// using the following formula.
//   acos(x) = -(i * log(x + i * sqrt((1 + x) * (1 - x))))
def : Pat<(CHLO_AcosOp NonComplexElementType:$input),
  (StableHLO_SelectOp
    (StableHLO_CompareOp
      $input,
      (ConstantLike<"-1"> $input),
      StableHLO_ComparisonDirectionValue<"NE">,
      (STABLEHLO_DEFAULT_COMPARISON_TYPE)
    ),
    (StableHLO_MulOp
      (ConstantLike<"2"> $input),
      (StableHLO_Atan2Op
        (StableHLO_SqrtOp
          (StableHLO_SubtractOp
            (ConstantLike<"1"> $input),
            (StableHLO_MulOp $input, $input)
          )
        ),
        (StableHLO_AddOp
          (ConstantLike<"1"> $input),
          $input
        )
      )
    ),
    (ConstantLike<"M_PI"> $input)
  )>;

// Expand acosh to MHLO dialect as follows:
//   acosh(x) = log(x + sqrt(x^2 - 1))      if x >= -1
//            = log(x + sqrt((x+1)*(x-1)))
//   acosh(x) = nan                         if x < -1
//
// If x^2 will overflow, we approximate sqrt(x^2 - 1) == x and compute as
// log(2*x) = log(2) + log(x).  (Note this works because negative x never
// overflows; x < -1 simply yields nan.
def : Pat<(CHLO_AcoshOp NonComplexElementType:$input),
  (StableHLO_SelectOp
    (StableHLO_CompareOp
      $input,
      (ConstantLike<"-1"> $input),
      StableHLO_ComparisonDirectionValue<"LT">,
      (STABLEHLO_DEFAULT_COMPARISON_TYPE)
    ),
    (ConstantLike<"NAN"> $input),
    (StableHLO_SelectOp
      (StableHLO_CompareOp
        $input,
        (StableHLO_SqrtOp
          (ConstantLikeMaxFiniteValue $input)
        ),
        StableHLO_ComparisonDirectionValue<"GE">,
        (STABLEHLO_DEFAULT_COMPARISON_TYPE)
      ),
      (StableHLO_AddOp
        (StableHLO_LogOp $input),
        (StableHLO_LogOp
          (ConstantLike<"2"> $input)
        )
      ),
      (StableHLO_LogOp
        (StableHLO_AddOp
          $input,
          (StableHLO_SqrtOp
            (StableHLO_MulOp
              (StableHLO_AddOp
                (ConstantLike<"1"> $input),
                $input
              ),
              (StableHLO_AddOp
                (ConstantLike<"-1"> $input),
                $input
              )
            )
          )
        )
      )
    )
  )>;

// Expand acosh for complex arguments to MHLO dialect as
//   acosh(x) = log(x + sqrt((x+1)*(x-1)))
//
// Per tensorflow/compiler/xla/client/lib/math.cc at the time of writing:
// "For now, we ignore the question of overflow if x is a
// complex type, because we don't yet have exhaustive tests for complex trig
// functions".
def : Pat<(CHLO_AcoshOp ComplexElementType:$input),
  (StableHLO_LogOp
    (StableHLO_AddOp
      $input,
      (StableHLO_SqrtOp
        (StableHLO_MulOp
          (StableHLO_AddOp
            $input,
            (ConstantLike<"1"> $input)
          ),
          (StableHLO_SubtractOp
            $input,
            (ConstantLike<"1"> $input)
          )
        )
      )
    )
  )>;


// Expand asin to MHLO dialect as follows:
//   asin(x) = 2 * atan(x / (1 + sqrt(1 - x^2)))
def : Pat<(CHLO_AsinOp $input),
  (StableHLO_MulOp
    (ConstantLike<"2"> $input),
    (StableHLO_Atan2Op
      $input,
      (StableHLO_AddOp
        (ConstantLike<"1"> $input),
        (StableHLO_SqrtOp
          (StableHLO_SubtractOp
            (ConstantLike<"1"> $input),
            (StableHLO_MulOp $input, $input)
          )
        )
      )
    )
  )>;

// Expand asinh for non-complex arguments to MHLO dialect as
//   asinh(x) = log(x + sqrt(x^2 + 1))
//
// If x^2 will overflow and x is positive, we can approximate x + sqrt(x^2 + 1)
// as 2*x and return log(2) + log(x).
//
// For small x, sqrt(x^2 + 1) will evaluate to 1 due to floating point
// arithmetic. However, we would like to retain the low order term of this,
// which is around 0.5 * x^2 using a binomial expansion.
// Let z = sqrt(a^2 + 1)
// The following rewrite retains the lower order term.
// log(a + sqrt(a^2 + 1))
//   = log((a + sqrt(a^2 + 1)) * (1 + sqrt(a^2 + 1)) / (1 + sqrt(a^2 + 1)))
//   = log((a + a^2 + 1 + a * z + z) / (1 + z))
//   = log(1 + a + a^2 / (1 + z))
//   = log(1 + a + a^2 / (1 + sqrt(a^2 + 1)))
//
// If x is negative, the above would give us some trouble; we can't approximate
// the result as x + abs(x) = 0 but we are saved by the fact that asinh(-x) =
// -asinh(x).
def : Pat<(CHLO_AsinhOp NonComplexElementType:$input),
  (StableHLO_MulOp
    (StableHLO_SignOp $input),
    (StableHLO_SelectOp
      (StableHLO_CompareOp
        (StableHLO_AbsOp $input),
        (StableHLO_SqrtOp
          (ConstantLikeMaxFiniteValue $input)
        ),
        StableHLO_ComparisonDirectionValue<"GE">,
        (STABLEHLO_DEFAULT_COMPARISON_TYPE)
      ),
      (StableHLO_AddOp
        (StableHLO_LogOp
          (StableHLO_AbsOp $input)
        ),
        (StableHLO_LogOp
          (ConstantLike<"2"> $input)
        )
      ),
      (StableHLO_SelectOp
        (StableHLO_CompareOp
          (StableHLO_AbsOp $input),
          (ConstantLike<"1"> $input),
          StableHLO_ComparisonDirectionValue<"LE">,
          (STABLEHLO_DEFAULT_COMPARISON_TYPE)
        ),
        (StableHLO_Log1pOp
          (StableHLO_AddOp
            (StableHLO_AbsOp $input),
            (StableHLO_MulOp
              (StableHLO_AbsOp $input),
              (StableHLO_DivOp
                (StableHLO_AbsOp $input),
                (StableHLO_AddOp
                  (ConstantLike<"1"> $input),
                  (StableHLO_SqrtOp
                    (StableHLO_AddOp
                      (StableHLO_MulOp
                        (StableHLO_AbsOp $input),
                        (StableHLO_AbsOp $input)
                      ),
                      (ConstantLike<"1"> $input)
                    )
                  )
                )
              )
            )
          )
        ),
        (StableHLO_LogOp
          (StableHLO_AddOp
            (StableHLO_AbsOp $input),
            (StableHLO_SqrtOp
              (StableHLO_AddOp
                (StableHLO_MulOp
                  (StableHLO_AbsOp $input),
                  (StableHLO_AbsOp $input)
                ),
                (ConstantLike<"1"> $input)
              )
            )
          )
        )
      )
    )
  )>;

// Expand asinh for complex arguments to MHLO dialect as
//   asinh(x) = log(x + sqrt(x^2 + 1))
//
// Per tensorflow/compiler/xla/client/lib/math.cc at the time of writing:
// "For now, we ignore the question of overflow if x is a
// complex type, because we don't yet have exhaustive tests for complex trig
// functions".
def : Pat<(CHLO_AsinhOp ComplexElementType:$input),
  (StableHLO_LogOp
    (StableHLO_AddOp
      $input,
      (StableHLO_SqrtOp
        (StableHLO_AddOp
          (StableHLO_MulOp $input, $input),
          (ConstantLike<"1"> $input)
        )
      )
    )
  )>;

// Express `atan` as
//   atan(x) = atan2(x, 1)
def : Pat<(CHLO_AtanOp $input),
  (StableHLO_Atan2Op
    $input,
    (ConstantLike<"1"> $input)
  )>;

// Express `atanh` for non-complex arguments as follows:
//   atanh(x) = 0.5 * log((1 + x) / (1 - x)) if abs(x) <= 1
//   atanh(x) = nan                          otherwise
def : Pat<(CHLO_AtanhOp NonComplexElementType:$input),
  (StableHLO_SelectOp
    (StableHLO_CompareOp
      (StableHLO_AbsOp $input),
      (ConstantLike<"1"> $input),
      StableHLO_ComparisonDirectionValue<"GT">,
      (STABLEHLO_DEFAULT_COMPARISON_TYPE)
    ),
    (ConstantLike<"NAN"> $input),
    (StableHLO_MulOp
      (StableHLO_SubtractOp
        (StableHLO_Log1pOp $input),
        (StableHLO_Log1pOp
          (StableHLO_NegOp $input)
        )
      ),
      (ConstantLike<"0.5"> $input)
    )
  )>;

// Express `atanh` for complex arguments as follows:
//   atanh(x) = (log(1 + x) - log(1 + (-x))) * 0.5
//
// Per tensorflow/compiler/xla/client/lib/math.cc at the time of writing:
// "For now, we ignore the nan edge case for complex inputs,
// because we don't yet have exhaustive tests for complex trig functions".
def : Pat<(CHLO_AtanhOp ComplexElementType:$input),
  (StableHLO_MulOp
    (StableHLO_SubtractOp
      (StableHLO_Log1pOp $input),
      (StableHLO_Log1pOp
        (StableHLO_NegOp $input)
      )
    ),
    (ConstantLike<"0.5"> $input)
  )>;

// Express `conj` as
//   conj(x) = (re(x), -im(x)).
def : Pat<(CHLO_ConjOp $v),
          (StableHLO_ComplexOp (StableHLO_RealOp $v), (StableHLO_NegOp (StableHLO_ImagOp $v)))>;

// Express `is_inf` as
//   is_inf(x) = is_pos_inf(|x|)
def : Pat<(CHLO_IsInfOp NonComplexElementType:$input),
  (CHLO_IsPosInfOp
    (StableHLO_AbsOp $input)
  )>;

// Express `is_pos_inf` as
//   is_pos_inf(x) = (x == +inf)
def : Pat<(CHLO_IsPosInfOp NonComplexElementType:$input),
  (StableHLO_CompareOp
    $input,
    (ConstantLikePosInfValue $input),
    StableHLO_ComparisonDirectionValue<"EQ">,
    (STABLEHLO_DEFAULT_COMPARISON_TYPE)
  )>;

// Express `is_neg_inf` as
//   is_neg_inf(x) = (x == -inf)
def : Pat<(CHLO_IsNegInfOp NonComplexElementType:$input),
  (StableHLO_CompareOp
    $input,
    (ConstantLikeNegInfValue $input),
    StableHLO_ComparisonDirectionValue<"EQ">,
    (STABLEHLO_DEFAULT_COMPARISON_TYPE)
  )>;

// Express `tan` as
//   sine(x) / cosine(x)
def : Pat<(CHLO_TanOp NonComplexElementType:$input),
  (StableHLO_DivOp
    (StableHLO_SineOp $input),
    (StableHLO_CosineOp $input)
  )>;


// Express `tan(a + bi)` as
//   (tan(a) + i tanh(b)) / (1 - i tan(a) * tanh(b))
def : Pat<(CHLO_TanOp ComplexElementType:$input),
  (StableHLO_DivOp
    (StableHLO_ComplexOp
      (CHLO_TanOp:$tan (StableHLO_RealOp $input)),
      (StableHLO_TanhOp:$tanh (StableHLO_ImagOp $input))),
    (StableHLO_ComplexOp
      (ConstantLike<"1.0"> $tan),
      (StableHLO_NegOp (StableHLO_MulOp $tan, $tanh)))
  )>;