einsum.hh

#ifndef __DUNE_ACFEM_TENSORS_OPTIMIZATION_EINSUM_HH__
#define __DUNE_ACFEM_TENSORS_OPTIMIZATION_EINSUM_HH__
 
#include "../../expressions/optimizegeneral.hh"
#include "../expressionoperations.hh"
#include "../operations/einsum.hh"
#include "../operationtraits.hh"
#include "../operations/transposetraits.hh"
#include "../modules.hh"
#include "policy.hh"
 
namespace Dune {
 
  namespace ACFem {
 
    template<>
    constexpr inline bool MultiplicationAdmitsScalarsV<EinsumOperation<Tensor::Seq<>, Tensor::Seq<>, Tensor::Seq<> > > = true;
 
    namespace Tensor::Optimization::Einsum {
 
      using DefaultTag = Tensor::Policy::DefaultOptimizationTag;
      using CommuteTag = OptimizeNext<DefaultTag>;
      using EinsumTag = OptimizeNext<CommuteTag, 2>;
 
 
      template<class F, class T0, class T1, class SFINAE = void>
      constexpr inline bool IsScalarEinsumV = false;
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      constexpr inline bool IsScalarEinsumV<
        OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >,
        T0, T1,
        std::enable_if_t<(EinsumRank<T0, Seq0, T1, Seq1> != 0
                          || !IsConstantExprArg<T0>::value
                          || !IsConstantExprArg<T1>::value
        )> > = false;
 
      template<class Op, class T0, class T1>
      using ScalarEinsumResult = std::decay_t<decltype(operate(DontOptimize{}, OperationTraits<Op>{}, std::declval<T0>(), std::declval<T1>())(Seq<>{}))>;
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      constexpr inline bool IsScalarEinsumV<
        OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >,
        T0, T1> = (!(ExamineOr<T0, IsRuntimeEqual>::value || ExamineOr<T1, IsRuntimeEqual>::value)
                   || IsFractionConstant<ScalarEinsumResult<EinsumOperation<Seq0, Seq1, Dims>, T0, T1> >::value);
 
      template<class F, class T0, class T1, std::enable_if_t<IsScalarEinsumV<F, T0 , T1>, int> = 0>
      constexpr decltype(auto) operate(OptimizeTerminal1, F&& f, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          tensor(std::move(operate(DontOptimize{}, f, std::forward<T0>(t0), std::forward<T1>(t1))(Seq<>{})))
          , "Compute Total Contractions"
          );
      }
 
 
      template<class F, class T0, class T1>
      constexpr inline bool ContractionOfFractionConstantsV = false;
 
      template<class Seq0, class Seq1, class Dims,
               class I0, I0 N0, I0 D0, class Signature0,
               class I1, I1 N1, I1 D1, class Signature1>
      constexpr inline bool ContractionOfFractionConstantsV<
        OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >,
        ConstantTensor<FractionConstant<I0, N0, D0>, Signature0>,
        ConstantTensor<FractionConstant<I1, N1, D1>, Signature1> > = EinsumRank<Signature0, Seq0, Signature1, Seq1> > 0;
 
      template<class F, class T0, class T1, std::enable_if_t<ContractionOfFractionConstantsV<F, T0, T1>, int> = 0>
      constexpr auto operate(OptimizeTerminal1, F&& f, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        using Signature = typename F::template Signature<T0, T1>;
        constexpr auto contractionDim = intFraction<F::contractionDimension()>();
        constexpr auto value = contractionDim * T0::data() * T1::data();
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          constantTensor(value, Signature{}, Disclosure{})
          , "Compile Time Constant Contraction"
          );
      }
 
 
      template<class T>
      constexpr inline std::size_t constness()
      {
        return (10*(std::size_t)IsConstantExprArg<T>::value
                + 100*(std::size_t)IsRuntimeEqualExpression<T>::value
                + 1000*(std::size_t)IsTypedValue<T>::value);
      }
 
      template<class T>
      constexpr std::size_t leftAffinity()
      {
        constexpr std::size_t byKind = constness<T>();
        constexpr std::size_t byWeight = (~0UL >> 16) - Expressions::weight<T>();
 
        return (byKind << 48) + byWeight;
      }
 
      template<class F, class T0, class T1>
      constexpr inline bool ShouldCommuteContractionV = false;
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      constexpr inline bool ShouldCommuteContractionV<OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >, T0, T1> =
        (Dims::size() > 0
         && EinsumRank<T0, Seq0, T1, Seq1> == 0
         && leftAffinity<T0>() < leftAffinity<T1>());
 
      template<class F, class T0, class T1, std::enable_if_t<ShouldCommuteContractionV<F, T0, T1>, int> = 0>
      constexpr decltype(auto) operate(CommuteTag, F&& f, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          (operate(std::forward<F>(f), std::forward<T1>(t1), std::forward<T0>(t0)))
          , "t0: " + t0.name() + "; t1: " + t1.name() + "; non scalar left affinity"
          );
      }
 
 
      template<class F, class T0, class T1>
      constexpr inline bool IsContractionOfConstWithEyeV = false;
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      constexpr inline bool IsContractionOfConstWithEyeV<OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >, T0, T1> =
        (Dims::size() > 0 // identify non-trivial contraction
         && IsEye<T0>::value
         && IsConstantTensor<T1>::value);
 
      template<class F, class T0, class T1, std::enable_if_t<IsContractionOfConstWithEyeV<F, T0, T1>, int> = 0>
      constexpr auto operate(DefaultTag, F&&, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        using EyeArg = std::decay_t<T0>;
        using EyeRest = SequenceSliceComplement<typename EyeArg::Signature, typename F::LeftIndexPositions>;
        using Const = std::decay_t<T1>;
        using ConstRest = SequenceSliceComplement<typename Const::Signature, typename F::RightIndexPositions>;
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          operate(ScalarEinsumFunctor{},
                  eye(EyeRest{}, Disclosure{}),
                  constantTensor(std::forward<T1>(t1).data(), ConstRest{}, Disclosure{}))
          , "einsum eye-ones"
          );
      }
 
      template<class F, class T0, class T1, std::enable_if_t<IsContractionOfConstWithEyeV<F, T1, T0>, int> = 0>
      constexpr auto operate(DefaultTag, F&& f, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        using Const = std::decay_t<T0>;
        using ConstRest = SequenceSliceComplement<typename Const::Signature, typename F::LeftIndexPositions>;
        using EyeArg = std::decay_t<T1>;
        using EyeRest = SequenceSliceComplement<typename EyeArg::Signature, typename F::RightIndexPositions>;
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          operate(ScalarEinsumFunctor{},
                  constantTensor(std::forward<T0>(t0).data(), ConstRest{}, Disclosure{}),
                  eye(EyeRest{}, Disclosure{}))
          , "einsum ones-eye"
          );
      }
 
 
    } // Tensors::Optimization::Einsum::
 
 
    namespace Expressions {
 
 
      template<std::size_t Pos0, std::size_t Pos1, std::size_t D0, class Field, class T0>
      constexpr inline bool ReturnFirstV<
        OperationTraits<EinsumOperation<IndexSequence<Pos0>, IndexSequence<Pos1>, IndexSequence<D0> > >,
        T0, Tensor::Eye<IndexSequence<D0, D0>, Field>
        > = Tensor::IsSelfTransposed<Transposition<Pos0, TensorTraits<T0>::rank-1>, T0>::value;
 
      template<std::size_t Pos0, std::size_t Pos1, std::size_t D0, class Field, class T1>
      constexpr inline bool ReturnSecondV<
        OperationTraits<EinsumOperation<IndexSequence<Pos0>, IndexSequence<Pos1>, IndexSequence<D0> > >,
        Tensor::Eye<IndexSequence<D0, D0>, Field>, T1,
        std::enable_if_t<!std::is_same<std::decay_t<T1>, Tensor::Eye<IndexSequence<D0, D0>, Field> >::value>
        > = Tensor::IsSelfTransposed<Transposition<0, Pos1>, T1>::value;
 
      template<class Seq0, class Seq1, class BlockSignature, class Field, class T0>
      constexpr inline bool ReturnFirstV<
        OperationTraits<EinsumOperation<Seq0, Seq1, BlockSignature> >,
        T0, Tensor::BlockEye<2, BlockSignature, Field>,
        std::enable_if_t<(BlockSignature::size() > 0
                          && isAlignedBlock<BlockSignature::size()>(Seq0{})
        )> > = (isAlignedBlock<BlockSignature::size()>(Seq1{})
                && Tensor::IsSelfTransposed<BlockTransposition<Seq0, MakeIndexSequence<BlockSignature::size(), TensorTraits<T0>::rank - BlockSignature::size()> >, T0>::value);
 
      template<class Seq0, class Seq1, class BlockSignature, class Field, class T1>
      constexpr inline bool ReturnSecondV<
        OperationTraits<EinsumOperation<Seq0, Seq1, BlockSignature> >,
        Tensor::BlockEye<2, BlockSignature, Field>, T1,
        std::enable_if_t<(BlockSignature::size() > 0
                          && isAlignedBlock<BlockSignature::size()>(Seq1{})
                          && !std::is_same<T1, Tensor::BlockEye<2, BlockSignature, Field> >::value
        )> > = (isAlignedBlock<BlockSignature::size()>(Seq0{})
                && Tensor::IsSelfTransposed<BlockTransposition<MakeIndexSequence<BlockSignature::size()>, Seq1>, T1>::value);
 
    } // Expressions::
 
 
    namespace Tensor::Optimization::Einsum {
 
      template<class F, class T0, class T1>
      constexpr inline bool IsContractionOfKroneckersV = false;
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      constexpr inline bool IsContractionOfKroneckersV<
        OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >,
        T0, T1
        > = (IsConstKroneckerDelta<T0>::value
             && IsConstKroneckerDelta<T1>::value);
 
      template<class F, class T0, class T1, std::enable_if_t<IsContractionOfKroneckersV<F, T0, T1>, int> = 0>
      constexpr decltype(auto) operate(OptimizeTerminal0, F&&, T0&&, T1&&)
      {
        using LeftPivots = typename T0::PivotSequence;
        using RightPivots = typename T1::PivotSequence;
        using LeftPos = typename F::LeftIndexPositions;
        using RightPos = typename F::RightIndexPositions;
        using Signature = typename F::template Signature<T0, T1>;
        if constexpr (std::is_same<SequenceSlice<LeftPivots, LeftPos>, SequenceSlice<RightPivots, RightPos> >::value) {
          DUNE_ACFEM_RECORD_OPTIMIZATION;
 
          using LPivot = SequenceSliceComplement<LeftPivots, LeftPos>;
          using RPivot = SequenceSliceComplement<RightPivots, RightPos>;
          using PivotIndices = SequenceCat<LPivot, RPivot>;
          DUNE_ACFEM_EXPRESSION_RESULT(
            kroneckerDelta(Signature{}, PivotIndices{}, Disclosure{})
            , "non-zero kronecker contraction"
            );
        } else {
          DUNE_ACFEM_RECORD_OPTIMIZATION;
 
          DUNE_ACFEM_EXPRESSION_RESULT(
            zeros(Signature{}, Disclosure{})
            , "zero kronecker contraction"
            );
        }
      }
 
 
      template<class F, class T0, class T1, class SFINAE = void>
      struct IsMiddleScalarNestedEinsum
        : FalseType
      {};
 
      template<class Seq0, class Seq1, class Dims, class T0, class T1>
      struct IsMiddleScalarNestedEinsum<
        OperationTraits<EinsumOperation<Seq0, Seq1, Dims> >, T0, T1,
        std::enable_if_t<(// identify nested multiplication
                          IsEinsumExpression<T1>::value
                          // identify non-scalars
                          && (TensorTraits<T0>::rank > 0)
                          && (TensorTraits<T1>::rank > 0)
                          // identify middel scalar
                          && TensorTraits<Operand<0, T1> >::rank == 0
        )> >
        : TrueType
      {};
 
      template<
        class F, class T0, class T1,
        std::enable_if_t<IsMiddleScalarNestedEinsum<F, T0, T1>::value, int> = 0>
      constexpr auto operate(OptimizeNext<EinsumTag>, F&& f, T0&& t0, T1&& t1)
      {
        DUNE_ACFEM_RECORD_OPTIMIZATION;
 
        DUNE_ACFEM_EXPRESSION_RESULT(
          operate(ScalarEinsumFunctor{},
                  std::forward<T1>(t1).operand(0_c),
                  operate(
                    std::forward<F>(f),
                    std::forward<T0>(t0),
                    std::forward<T1>(t1).operand(1_c)
                    )
            )
          , "t0: " + t0.name() + "; t1: " + t1.name() + "; middle scalar einsum"
          );
      }
 
 
 
    } // Tensor::Optimization::Einsum::
 
    namespace Expressions {
 
      using Tensor::Optimization::Einsum::operate;
 
    } // NS Expressions
 
  } // NS ACFem
 
} // NS Dune
 
#endif // __DUNE_ACFEM_TENSORS_OPTIMIZATION_MULTIPY_HH__