Further tuning on flattenings of rings (#4400)

* Attempt to speed up flattenings once more.

* Extend speedup and make caching of elements an option.

src/Rings/MPolyMap/flattening_specializations.jl

@@ -0,0 +1,59 @@
+# specialized routines with improvements
+
+# some unifying getters to allow for generic code
+_poly_lift(x::MPolyRingElem) = x
+_poly_lift(x::MPolyQuoRingElem) = x.f
+
+_poly_ring(R::MPolyRing) = R
+_poly_ring(R::MPolyQuoRing) = base_ring(R)
+
+function (phi::RingFlattening{<:Union{
+                                    <:MPolyRing{<:MPolyRingElem},
+                                    <:MPolyRing{<:MPolyQuoRingElem},
+                                    <:MPolyQuoRing{<:MPolyRingElem{<:MPolyRingElem}},
+                                    <:MPolyQuoRing{<:MPolyRingElem{<:MPolyQuoRingElem}}
+                                   }, 
+                              <:Union{
+                                    <:MPolyRing, <:MPolyQuoRing
+                                   }})(x::RingElem)
+  !phi.cached && return _compute_flattened_element(phi, x)
+  return get!(flat_counterparts(phi), x) do
+    _compute_flattened_element(phi, x)
+  end::elem_type(codomain(phi))
+end
+
+function _compute_flattened_element(phi, x)
+  is_zero(x) && return zero(codomain(phi))
+  R = domain(phi)
+  parent(x) === R || return _compute_flattened_element(phi, R(x))
+  RP = _poly_ring(R)
+  P = coefficient_ring(RP)
+  PP = _poly_ring(P)
+  kk = coefficient_ring(PP)
+  S = codomain(phi)
+  SP = _poly_ring(S)
+  @assert kk === coefficient_ring(SP)
+  if is_constant(_poly_lift(x))
+    c = first(AbstractAlgebra.coefficients(_poly_lift(x)))
+    is_constant(_poly_lift(c)) && return S(first(AbstractAlgebra.coefficients(_poly_lift(c))))
+  end
+  res_ctx = MPolyBuildCtx(SP)
+  for (c, e) in zip(AbstractAlgebra.coefficients(_poly_lift(x)), AbstractAlgebra.exponent_vectors(_poly_lift(x)))
+    for (cc, ee) in zip(AbstractAlgebra.coefficients(_poly_lift(c)), AbstractAlgebra.exponent_vectors(_poly_lift(c)))
+      push_term!(res_ctx, cc, vcat(e, ee))
+    end
+  end
+  return S(finish(res_ctx))
+end
+
+# TODO: One can still try to do something like the above optimizations 
+# also for polynomial rings over localizations of polynomial rings. The 
+# problem there is that one can not directly concatenate exponent vectors 
+# because we have potential denominators floating around. But should this 
+# turn out to be a bottleneck, one can probably think of something.
+
+# TODO: Think of a way to speed up the inverse of a RingFlattening.
+# The problem is that at the moment the inverse is a standalone map, 
+# stored in the internals of a RingFlattening. When it's applied to an 
+# element, it does not know about its origin anymore, so we can not 
+# overwrite this via some dispatch for RingFlattening. 

src/Rings/MPolyMap/flattenings.jl

@@ -19,11 +19,16 @@
   orig_to_flat::MPolyAnyMap{TowerRingType, FlatRingType}
   flat_to_orig::Map{FlatRingType, TowerRingType}
   base_ring_to_flat::Map{CoeffRingType, FlatRingType}
+  cached::Bool # This is an internal variable to turn on/off caching of 
+               # counterparts of ring elements. It is turned off by default, 
+               # because one has to be careful with in-place operations.
+               # Note that parent-like objects are still always cached!
 
   flat_counterparts::AbstractAlgebra.WeakKeyIdDict{Any, Any}
 
   function RingFlattening(
-      S::MPolyRing{RingElemType}
+      S::MPolyRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem}
     R = base_ring(S)
     kk = coefficient_ring(R)
@@ -33,12 +38,13 @@
     S_flat_to_S = hom(S_flat, S, vcat(gens(S), S.(gens(R))), check=false)
     return new{typeof(S), typeof(S_flat), typeof(R)}(S, S_flat, R, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     R_to_S_flat
+                                                     R_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyDecRing{RingElemType}
+      S::MPolyDecRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem}
     R = base_ring(S)
     kk = coefficient_ring(R)
@@ -57,12 +63,13 @@
     S_flat_to_S = hom(S_flat, S, vcat(gens(S), S.(gens(R))), check=false)
     return new{typeof(S), typeof(S_flat), typeof(R)}(S, S_flat, R, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     R_to_S_flat
+                                                     R_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyRing{RingElemType}
+      S::MPolyRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyQuoRingElem}
     A = base_ring(S)::MPolyQuoRing
     R = base_ring(A)::MPolyRing
@@ -84,12 +91,13 @@
 
     return new{typeof(S), typeof(S_flat), typeof(A)}(S, S_flat, A, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     A_to_S_flat
+                                                     A_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyDecRing{RingElemType}
+      S::MPolyDecRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyQuoRingElem}
     A = base_ring(S)::MPolyQuoRing
     R = base_ring(A)::MPolyRing
@@ -120,12 +128,13 @@
 
     return new{typeof(S), typeof(S_flat), typeof(A)}(S, S_flat, A, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     A_to_S_flat
+                                                     A_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyRing{RingElemType}
+      S::MPolyRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyQuoLocRingElem}
     Q = base_ring(S)::MPolyQuoLocRing
     R = base_ring(Q)::MPolyRing
@@ -153,12 +162,13 @@
 
     return new{typeof(S), typeof(S_flat), typeof(Q)}(S, S_flat, Q, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     Q_to_S_flat
+                                                     Q_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyDecRing{RingElemType}
+      S::MPolyDecRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyQuoLocRingElem}
     Q = base_ring(S)::MPolyQuoLocRing
     R = base_ring(Q)::MPolyRing
@@ -201,12 +211,13 @@
 
     return new{typeof(S), typeof(S_flat), typeof(Q)}(S, S_flat, Q, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     Q_to_S_flat
+                                                     Q_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyRing{RingElemType}
+      S::MPolyRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyLocRingElem}
     L = base_ring(S)::MPolyLocRing
     R = base_ring(L)::MPolyRing
@@ -228,12 +239,13 @@
 
     return new{typeof(S), typeof(S_flat), typeof(L)}(S, S_flat, L, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     L_to_S_flat
+                                                     L_to_S_flat, cached
                                                     )
   end
 
   function RingFlattening(
-      S::MPolyDecRing{RingElemType}
+      S::MPolyDecRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyLocRingElem}
     L = base_ring(S)::MPolyLocRing
     R = base_ring(L)::MPolyRing
@@ -265,14 +277,15 @@
 
     return new{typeof(S), typeof(S_flat), typeof(L)}(S, S_flat, L, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     L_to_S_flat
+                                                     L_to_S_flat, cached
                                                     )
   end
 
 
   # Flattenings of quotient rings of the form (𝕜[x][u])/J → 𝕜[x, u]/J'
   function RingFlattening(
-      S::MPolyQuoRing{RingElemType}
+      S::MPolyQuoRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem{<:MPolyRingElem}}
     P = base_ring(S) # the free polynomial ring
     P_flattening = flatten(P)
@@ -287,13 +300,14 @@
     S_flat_to_S = hom(S_flat, S, vcat(gens(S), S.(gens(R))), check=false)
     return new{typeof(S), typeof(S_flat), typeof(R)}(S, S_flat, R, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     R_to_S_flat
+                                                     R_to_S_flat, cached
                                                     )
   end
 
   # Flattenings of quotient rings of the form ((𝕜[x]/I)[u])/J → 𝕜[x, u]/(I' + J')
   function RingFlattening(
-      S::MPolyQuoRing{RingElemType}
+      S::MPolyQuoRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem{<:MPolyQuoRingElem}}
     P = base_ring(S)::MPolyRing # the polynomial ring behind S
     A = base_ring(P)::MPolyQuoRing # the coefficient ring of S
@@ -310,13 +324,14 @@
 
     return new{typeof(S), typeof(S_flat), typeof(A)}(S, S_flat, A, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     A_to_S_flat
+                                                     A_to_S_flat, cached
                                                     )
   end
 
   # Flattenings of quotient rings of the form (((𝕜[x]/I)[U⁻¹])[u])/J → (𝕜[x, u]/(I' + J'))[U'⁻¹]
   function RingFlattening(
-      S::MPolyQuoRing{RingElemType}
+      S::MPolyQuoRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem{<:MPolyQuoLocRingElem}}
     P = base_ring(S)::MPolyRing
     Q = base_ring(P)::MPolyQuoLocRing
@@ -349,13 +364,14 @@
 
     return new{typeof(S), typeof(S_flat), typeof(Q)}(S, S_flat, Q, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     Q_to_S_flat
+                                                     Q_to_S_flat, cached
                                                     )
   end
 
   # Flattenings of quotient rings of the form ((𝕜[x][U⁻¹])[u])/J → (𝕜[x, u])[U'⁻¹]/J'
   function RingFlattening(
-      S::MPolyQuoRing{RingElemType}
+      S::MPolyQuoRing{RingElemType};
+      cached::Bool=false
     ) where {RingElemType <: MPolyRingElem{<:MPolyLocRingElem}}
     P = base_ring(S)::MPolyRing
     L = base_ring(P)::MPolyLocRing
@@ -385,19 +401,17 @@
     S_flat_to_S = MapFromFunc(S_flat, S, my_map)
     return new{typeof(S), typeof(S_flat), typeof(L)}(S, S_flat, L, 
                                                      S_to_S_flat, S_flat_to_S,
-                                                     L_to_S_flat
+                                                     L_to_S_flat, cached
                                                     )
   end
 end
 
 ### Getters 
 function (phi::RingFlattening)(x::RingElem)
-  if !haskey(flat_counterparts(phi), x)
-    result = phi.orig_to_flat(x)
-    flat_counterparts(phi)[x] = result
-    return result
-  end
- return flat_counterparts(phi)[x]::elem_type(codomain(phi))
+  !phi.cached && return phi.orig_to_flat(x)
+  return get!(flat_counterparts(phi), x) do
+    phi.orig_to_flat(x)
+  end::elem_type(codomain(phi))
 end
 
 function inverse(phi::RingFlattening)
@@ -441,12 +455,16 @@ function flat_counterparts(phi::RingFlattening)
 end
 
 ### Some basic functionality
-@attr RingFlattening function flatten(R::MPolyRing)
-  return RingFlattening(R)
+function flatten(R::MPolyRing; cached::Bool=false)
+  return get_attribute!(R, :flatten) do
+    RingFlattening(R; cached)
+  end::RingFlattening
 end
 
-@attr RingFlattening function flatten(R::MPolyQuoRing)
-  return RingFlattening(R)
+function flatten(R::MPolyQuoRing; cached::Bool=false)
+  return get_attribute!(R, :flatten) do
+    RingFlattening(R; cached)
+  end::RingFlattening
 end
 
 function (phi::RingFlattening)(I::MPolyIdeal)

src/Rings/Rings.jl

@@ -26,6 +26,7 @@ include("mpoly-localization_types.jl")
 include("mpoly-localizations.jl")
 include("mpolyquo-localizations.jl")
 include("MPolyMap/flattenings.jl")
+include("MPolyMap/flattening_specializations.jl")
 include("FinField.jl")
 include("NumberField.jl")
 include("FunctionField.jl")

test/Rings/MPolyAnyMap/flattenings.jl

@@ -16,7 +16,7 @@
 
   R, _ = quo(R, ideal(R, sum(gens(R))))
   S, _ = polynomial_ring(R, [:s, :t])
-  phi = Oscar.flatten(S)
+  phi = Oscar.flatten(S; cached=true)
   @test vcat(phi.(gens(S)), Oscar.map_from_coefficient_ring_to_flattening(phi).(gens(coefficient_ring(S)))) == gens(codomain(phi))
   @test gens(S) == inverse(phi).(phi.(gens(S)))
   @test Oscar.map_from_coefficient_ring_to_flattening(phi).(gens(R)) == phi.(S.(gens(R)))
@@ -301,7 +301,8 @@ end
 @testset "garbage collection" begin
   R, (x,y,z) = polynomial_ring(QQ, [:x, :y, :z], cached=false)
   S, _ = polynomial_ring(R, [:s, :t], cached=false)
-  phi = Oscar.flatten(S)
+  S, _ = quo(S, ideal(S, elem_type(S)[]))
+  phi = Oscar.flatten(S; cached=true)
   # julia might not consider `x^2+y^2` as being unused until the end of the scope
   # see https://github.com/JuliaLang/julia/issues/51818
   # so we put this into a function