/* * Copyright (c) 2022-2024, NVIDIA CORPORATION. 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. */ #pragma once #pragma nv_diag_suppress 177, 550 #include #include "Cn.h" //---------------------------------------------------------------------------- // Rn: ranged integer (with size and multiplier) //---------------------------------------------------------------------------- enum Kind { NONE, UNROLL, ID, }; template class Ranged { static_assert(multiplier_); public: typedef decltype(size_) type; FT_DEV_CEXPR Ranged(type var_ = 0) : var(var_) { } static constexpr auto min = type(size_ && multiplier_ < 0 ? size_ - 1 : 0) * multiplier_; static constexpr auto max = type(size_ && multiplier_ > 0 ? size_ - 1 : 0) * multiplier_; static constexpr auto abs = max - min; static constexpr bool minInf = (size_ == 0 && multiplier_ < 0); static constexpr bool maxInf = (size_ == 0 && multiplier_ > 0); static constexpr bool inf = (size_ == 0); static constexpr auto zero = decltype(size_)(0); static constexpr auto ZERO = decltype(size_ * multiplier_)(0); static constexpr Kind kind = kind_; static constexpr type size = size_; static constexpr auto multiplier = multiplier_; type var; FT_DEV_CEXPR auto operator+() const { return *this; } FT_DEV_CEXPR auto operator-() const { return Ranged{var}; } }; template ? 0 : decltype(kind_)(1), auto multiplier_ = decltype(size_)(1)> using Rn = std::conditional_t, Ranged, Ranged>; //---------------------------------------------------------------------------- // Poly: polynomial integer //---------------------------------------------------------------------------- template class Poly { public: typedef std::tuple Terms; FT_DEV_CEXPR Poly(Cn, Terms ts_) : terms(ts_) { } FT_DEV_CEXPR Poly(Cn, Ts_... ts_) : terms(ts_...) { } FT_DEV_CEXPR Poly(Terms ts_) : terms(ts_) { } FT_DEV_CEXPR Poly(Ts_... ts_) : terms(ts_...) { } static constexpr auto min = (bias_ + ... + Ts_::min); static constexpr auto max = (bias_ + ... + Ts_::max); static constexpr bool minInf = (false || ... || Ts_::minInf); static constexpr bool maxInf = (false || ... || Ts_::maxInf); static constexpr auto zero = decltype(bias_)(0); static constexpr auto ZERO = decltype((bias_ + ... + Ts_::zero))(0); static constexpr auto bias = bias_; Terms terms; FT_DEV_CEXPR auto operator+() const { return *this; } FT_DEV_CEXPR auto operator-() const { return negateImp(std::index_sequence_for()); } template FT_DEV_CEXPR auto mul(Cn) const { return mulImp(cn, cn); } template FT_DEV_CEXPR auto operator/(Cn) const { return divImp(cn, cn); } template FT_DEV_CEXPR auto mod(Cn) const { return modImp(cn, cn); } template FT_DEV_CEXPR auto filter(Cn) const { return filterImp(cn, cn); } template FT_DEV_CEXPR auto filterDiv(Cn) const { if constexpr (b_ == 0) return *this; // return itself if indivisible else if constexpr (!divisible(cn)) return *this; // return itself if indivisible else return filterDivImp(cn, cn); } template FT_DEV_CEXPR static bool divisible(Cn) { static_assert(b_); constexpr auto dMin = divisibleMin(cn, std::index_sequence_for()); constexpr auto dMax = divisibleMax(cn, std::index_sequence_for()); constexpr auto iMin = indivisibleMin(cn, std::index_sequence_for()); constexpr auto iMax = indivisibleMax(cn, std::index_sequence_for()); return dMin == 0 && iMin == 0 && bias_ >= 0 && iMax + bias_ % b_ < cexpr_abs(b_) || dMax == 0 && iMax == 0 && bias_ <= 0 && iMin + bias_ % b_ > -cexpr_abs(b_); } template FT_DEV_CEXPR static bool hasOnly(Cn) { return hasOnlyImp(cn, cn); } private: template FT_DEV_CEXPR auto negateImp(std::index_sequence) const { return Poly<-bias_, decltype(-std::get(terms))...>{cn<-bias_>, std::tuple{-std::get(terms)...}}; } template FT_DEV_CEXPR static auto divisibleMin(Cn, std::index_sequence) { return (zero + ... + (std::tuple_element_t::multiplier % b_ && std::tuple_element_t::size != 1 ? std::tuple_element_t::ZERO : std::tuple_element_t::min + std::tuple_element_t::minInf * -cexpr_abs(b_))); } template FT_DEV_CEXPR static auto divisibleMax(Cn, std::index_sequence) { return (zero + ... + (std::tuple_element_t::multiplier % b_ && std::tuple_element_t::size != 1 ? std::tuple_element_t::ZERO : std::tuple_element_t::max + std::tuple_element_t::maxInf * cexpr_abs(b_))); } template FT_DEV_CEXPR static auto indivisibleMin(Cn, std::index_sequence) { return (zero + ... + (std::tuple_element_t::multiplier % b_ && std::tuple_element_t::size != 1 ? std::tuple_element_t::min + std::tuple_element_t::minInf * -cexpr_abs(b_) > -cexpr_abs(b_) ? std::tuple_element_t::min + std::tuple_element_t::minInf * -cexpr_abs(b_) : -cexpr_abs(b_) + std::gcd(cexpr_abs(b_), std::tuple_element_t::multiplier) : std::tuple_element_t::ZERO)); } template FT_DEV_CEXPR static auto indivisibleMax(Cn, std::index_sequence) { return (zero + ... + (std::tuple_element_t::multiplier % b_ && std::tuple_element_t::size != 1 ? std::tuple_element_t::max + std::tuple_element_t::maxInf * cexpr_abs(b_) < cexpr_abs(b_) ? std::tuple_element_t::max + std::tuple_element_t::maxInf * cexpr_abs(b_) : cexpr_abs(b_) - std::gcd(cexpr_abs(b_), std::tuple_element_t::multiplier) : std::tuple_element_t::ZERO)); } template FT_DEV_CEXPR static bool hasOnlyImp(Cn, Cn) { if constexpr (i_ == 0) return true; else if constexpr (std::tuple_element_t::kind != kind_ && std::tuple_element_t::size != 1) return false; else return hasOnlyImp(cn, cn); } template FT_DEV_CEXPR auto mulImp(Cn, Cn) const { if constexpr (i_ == 0) return cn; else return mulImp(cn, cn) + std::get(terms) * cn; } template FT_DEV_CEXPR auto divImp(Cn, Cn) const { static_assert(b_); static_assert(divisible(cn)); if constexpr (i_ == 0) return cn; else if constexpr (std::tuple_element_t::abs >= cexpr_abs(b_) || std::tuple_element_t::inf) return divImp(cn, cn) + std::get(terms) / cn; else return divImp(cn, cn); } template FT_DEV_CEXPR auto modImp(Cn, Cn) const { static_assert(b_); static_assert(divisible(cn)); if constexpr (i_ == 0) return cn; else if constexpr (std::tuple_element_t::multiplier % b_ && std::tuple_element_t::size != 1) return modImp(cn, cn) + std::get(terms) % cn; else return modImp(cn, cn); } template FT_DEV_CEXPR auto filterImp(Cn, Cn) const { if constexpr (i_ == 0) return Poly{}; else if constexpr (std::tuple_element_t::kind == kind_ && std::tuple_element_t::size != 1) return filterImp(cn, cn) + std::get(terms); else return filterImp(cn, cn); } template FT_DEV_CEXPR auto filterDivImp(Cn, Cn) const { static_assert(b_); static_assert(divisible(cn)); if constexpr (i_ == 0) return Poly{}; else if constexpr (std::tuple_element_t::abs >= cexpr_abs(b_) || std::tuple_element_t::inf) return filterDivImp(cn, cn) + std::get(terms); else return filterDivImp(cn, cn); } }; // constructs Poly from Cn and Rns template Poly(Cn, Ranged...) -> Poly...>; // constructs Poly from Rns template Poly(Ranged...) -> Poly...>; //---------------------------------------------------------------------------- // Operators for Rn and Poly //---------------------------------------------------------------------------- /* We should never use int * Rn template FT_DEV_CEXPR std::enable_if_t, Rn> operator * (T_ a_, Ranged x_) { return Rn{x_.var * a_}; } template FT_DEV_CEXPR std::enable_if_t, Rn> operator * (Ranged x_, T_ b_) { return Rn{x_.var * b_}; } */ template FT_DEV_CEXPR std::enable_if_t> operator*(Cn, Ranged x_) { return Rn{x_.var}; } template FT_DEV_CEXPR std::enable_if_t> operator*(Ranged x_, Cn) { return Rn{x_.var}; } template FT_DEV_CEXPR std::enable_if_t> operator/(Ranged x_, Cn) { return Rn{x_.var}; } template FT_DEV_CEXPR std::enable_if_t> operator%(Ranged x_, Cn) { return cn; } template FT_DEV_CEXPR Rn operator<<(Ranged x_, Cn) { return Rn{x_.var}; } template FT_DEV_CEXPR std::enable_if_t> b_)>> operator>>(Ranged x_, Cn) { return Rn> b_)>{x_.var}; } template FT_DEV_CEXPR std::enable_if_t<(Rn::abs < cexpr_abs(b_) && !Rn::inf && m_ % b_ != 0), Cn> operator/(Ranged x_, Cn) { return cn; } template FT_DEV_CEXPR std::enable_if_t<(Rn::abs < cexpr_abs(b_) && !Rn::inf && m_ % b_ != 0), Rn> operator%(Ranged x_, Cn) { return Rn{x_.var}; } template FT_DEV_CEXPR std::enable_if_t<(Rn::abs < (1 << b_) && !Rn::inf && m_ % (1 << b_) != 0), Cn<(m_ >> b_)>> operator>>(Ranged x_, Cn) { return cn<(m_ >> b_)>; } template FT_DEV_CEXPR std::enable_if_t<(Rn::abs >= cexpr_abs(b_) || Rn::inf) && (m_ % b_ != 0 && b_ % m_ == 0), Rn> operator/(Ranged x_, Cn) { return Rn{x_.var / cexpr_abs(b_ / m_)}; } template FT_DEV_CEXPR std::enable_if_t<(Rn::abs >= cexpr_abs(b_) || Rn::inf) && (m_ % b_ != 0 && b_ % m_ == 0), Rn> operator%(Ranged x_, Cn) { return Rn{x_.var % cexpr_abs(b_ / m_)}; } template FT_DEV_CEXPR std::enable_if_t<(m_ > 0) && // correct only when positive (Rn::abs >= (1 << b_) || Rn::inf) && (m_ % (1 << b_) != 0 && (1 << b_) % m_ == 0), Rn> operator>>(Ranged x_, Cn) { return Rn{ x_.var / cexpr_abs((1 << b_) / m_)}; } template FT_DEV_CEXPR auto operator+(Ranged a_, Cn) { return Poly{cn, std::tuple{a_}}; } template FT_DEV_CEXPR auto operator-(Ranged a_, Cn) { return Poly{cn<-b_>, std::tuple{a_}}; } template FT_DEV_CEXPR auto operator+(Cn, Ranged b_) { return Poly{cn, std::tuple{b_}}; } template FT_DEV_CEXPR auto operator-(Cn, Ranged b_) { return Poly{cn, std::tuple{-b_}}; } template FT_DEV_CEXPR auto operator+(Poly a_, Cn) { return Poly{cn, a_.terms}; } template FT_DEV_CEXPR auto operator-(Poly a_, Cn) { return Poly{cn, a_.terms}; } template FT_DEV_CEXPR auto operator+(Cn, Poly b_) { return Poly{cn, b_.terms}; } template FT_DEV_CEXPR auto operator-(Cn, Poly b_) { return Poly{cn, (-b_).terms}; } template FT_DEV_CEXPR auto operator+(Ranged a_, Ranged b_) { return Poly{cn, std::tuple{a_, b_}}; } template FT_DEV_CEXPR auto operator-(Ranged a_, Ranged b_) { return Poly{cn, std::tuple{a_, -b_}}; } template FT_DEV_CEXPR auto operator+(Poly a_, Ranged b_) { return Poly{cn, std::tuple_cat(a_.terms, std::tuple{b_})}; } template FT_DEV_CEXPR auto operator-(Poly a_, Ranged b_) { return Poly{cn, std::tuple_cat(a_.terms, std::tuple{-b_})}; } template FT_DEV_CEXPR auto operator+(Ranged a_, Poly b_) { return Poly{cn, std::tuple_cat(std::tuple{a_}, b_.terms)}; } template FT_DEV_CEXPR auto operator-(Ranged a_, Poly b_) { return Poly{cn<-B_>, std::tuple_cat(std::tuple{a_}, (-b_).terms)}; } template FT_DEV_CEXPR auto operator+(Poly a_, Poly b_) { return Poly{cn, std::tuple_cat(a_.terms, b_.terms)}; } template FT_DEV_CEXPR auto operator-(Poly a_, Poly b_) { return Poly{cn, std::tuple_cat(a_.terms, (-b_).terms)}; } template FT_DEV_CEXPR std::enable_if_t>().mul(cn))> operator*( Cn, Poly x_) { return x_.mul(cn); } template FT_DEV_CEXPR std::enable_if_t>().mul(cn))> operator*( Poly x_, Cn) { return x_.mul(cn); } template FT_DEV_CEXPR std::enable_if_t>().mod(cn))> operator%( Poly x_, Cn) { return x_.mod(cn); } template FT_DEV_CEXPR auto operator<<(Poly x_, Cn) { return x_ * cn<(1 << b_)>; } template FT_DEV_CEXPR auto operator>>(Poly x_, Cn) { return x_ / cn<(1 << b_)>; } template FT_DEV_CEXPR auto operator*(Ranged a_, Ranged b_) { return Ranged < kA == kB ? kA : NONE, zA * zB, mA * mB > {a_.var * b_.var}; } template FT_DEV_CEXPR auto operator*(Poly a_, Ranged b_) { return Ranged < Poly::hasOnly(cn) ? k_ : NONE > {get(a_) * b_.var} * cn; } template FT_DEV_CEXPR auto operator*(Ranged a_, Poly b_) { return Ranged < Poly::hasOnly(cn) ? k_ : NONE > {a_.var * get(b_)} * cn; } /* We should never use Poly * Poly template FT_DEV_CEXPR auto operator * (Poly a_, Poly b_) { return Ranged::hasOnly(cn) && Poly::hasOnly(cn) ? UNROLL : ( Poly::hasOnly(cn) && Poly::hasOnly(cn) ? ID : NONE)>{get(a_) * get(b_)}; } */ //---------------------------------------------------------------------------- // get() for Cn, Rn and Poly //---------------------------------------------------------------------------- template FT_DEV_CEXPR auto get(Cn) { return value_; } template FT_DEV_CEXPR auto get(Ranged x_) { if constexpr (size_ == 1) return Rn::ZERO; else return x_.var * multiplier_; } template FT_DEV_CEXPR auto getImp(Poly x_, std::index_sequence) { return (bias_ + ... + get(std::get(x_.terms))); } template FT_DEV_CEXPR auto get(Poly x_) { return getImp(x_.filter(cn), std::make_index_sequence))::Terms>>()); } template FT_DEV_CEXPR auto get(Poly x_) { return bias_ + get(x_) + get(x_) + get(x_); } //---------------------------------------------------------------------------- // Comparison operators for Rn and Poly //---------------------------------------------------------------------------- template FT_DEV_CEXPR auto ltzeroImp(Poly x_) { constexpr decltype(A_) p2 = A_ ? ((-A_) ^ (-A_ - 1)) / 2 + 1 : 0; constexpr auto n1 = std::tuple_size_v) .filterDiv(cn<1>))::Terms>; constexpr auto n2 = std::tuple_size_v) .filterDiv(cn))::Terms>; constexpr auto nA = std::tuple_size_v) .filterDiv(cn))::Terms>; if constexpr (Poly::min >= 0 && !Poly::minInf) return cn; else if constexpr (Poly::max < 0 && !Poly::maxInf) return cn; else if constexpr (A_ < 0 && nA < n2 && nA < n1) return ltzeroImp((x_ - cn) .filterDiv(cn) + cn); else if constexpr (A_ < 0 && n2 < n1) return ltzeroImp((x_ - cn) .filterDiv(cn) + cn); else if (A_ + get(x_) + decltype(x_.filter(cn) + x_.filter(cn))::min >= 0 && !decltype(x_.filter(cn) + x_.filter(cn))::minInf) return false; else if (A_ + get(x_) + decltype(x_.filter(cn) + x_.filter(cn))::max < 0 && !decltype(x_.filter(cn) + x_.filter(cn))::maxInf) return true; else if constexpr (A_ < 0) return get(x_) + get