Option to invalidate the array_cache for LazyArray

[tam724]

Would it make sense to have a way to invalidate the cache for indexing a LazyArray? So that essentially this check in getindex!(...)

Gridap.jl/src/Arrays/LazyArrays.jl

Line 199 in 55a399a

if index_and_item.index != index

is skipped and the getindex!(...) call always reevaluates the cache entries (even if the same item was indexed before).

As a MWE I would have liked something like this to work on the first getindex!(...) call:

using Gridap

model = CartesianDiscreteModel((0, 1), 10)
V = TestFESpace(model, ReferenceFE(lagrangian, Float64, 1), conformity=:H1)
Ω = Triangulation(model)
dx = Measure(Ω, 4)

α = Ref(1.0) # I would like to change this value and reevaluate the quadrature without having to build the whole lazy expression again.
g(x) = exp(-α[] * x[1])

b(v) = ∫(g * v)dx

v = get_fe_basis(V)
cell_vec_lazy = b(v)[Ω]
cell_vec_cache = Gridap.array_cache(cell_vec_lazy)
getindex!(cell_vec_cache, cell_vec_lazy, 1)[1]

Gridap.assemble_vector(b, V)[1] # this should return the same value as the getindex!(...) before

# now change α
α[] = 2.0
getindex!(cell_vec_cache, cell_vec_lazy, 1)[1] # returns the same value as before (where α still was 1)
getindex!(cell_vec_cache, cell_vec_lazy, 2)[1]; # this is a workaround to invalidate the cache (simply index into another cell)
getindex!(cell_vec_cache, cell_vec_lazy, 1)[1] # now this returns the value where α = 2

The cache invalidation (here by indexing another cell) could be more explicit and skip the reevaluation of the cache entries. For me implementing something along the lines of this was enough (but I only understand the type-structure of the cache to some extent):

invalidate_cache!(::Union{Nothing, Gridap.Arrays.CachedVector, Gridap.Arrays.CachedMatrix}) = nothing
function invalidate_cache!(cache::Tuple{T1}) where T1
    invalidate_cache!(cache[1])
end
function invalidate_cache!(cache::Tuple{T1, T2}) where {T1, T2}
    invalidate_cache!(cache[1])
    invalidate_cache!(cache[2])
end
function invalidate_cache!(cache::Tuple{T1, T2, T3}) where {T1, T2, T3}
    invalidate_cache!(cache[1])
    invalidate_cache!(cache[2])
    invalidate_cache!(cache[3])
end
function invalidate_cache!(cache::Tuple{T, <:Gridap.Arrays.IndexItemPair}) where T
    invalidate_cache!(cache[1])
    cache[2].index = -1
end

Alternatively, getindex!() could have an optional argument always_reevaluate=false that switches off the index check.

My use case

I stumbled upon this while I was experimenting with a probably rather uncommon use case for a FE-library. I wanted to assemble a matrix

a(u, v) = ∫( ∫( k(x₁, x₂) * u(x₁) )dx₁ * v(x₂) ) dx₂
However I thought that I could reuse the quadrature rules, basis function definitions and grid-definitions in Gridap.jl.

Using only high-level interfaces, I could implement the following:

k(x1, x2) = exp(-50.0*(x1-x2)^2)

model = CartesianDiscreteModel((0, 1), 10)
V = TestFESpace(model, ReferenceFE(lagrangian, Float64, 1), conformity=:H1)
Ω = Triangulation(model)
dx = Measure(Ω, 4)
K = zeros(num_free_dofs(V), num_free_dofs(V))
u = FEFunction(V, zeros(num_free_dofs(V)))
for i in 1:num_free_dofs(V)
    u.free_values[i] = 1.0
    K[i, :] = Gridap.assemble_vector(
        v -> ∫(
            (y -> sum(∫((x -> k(y[1], x[1])) * u)dx))
                * v)dx
        , V)
    u.free_values[i] = 0.0
end

However, the inner quadrature always evaluates the full domain whereas only the cells in the support of the basis u[i] should be integrated. For larger N that quickly becomes prohibitively slow.

Leveraging the low-level APIs of Gridap, I found the getindex! to be essentially the function I needed for cellwise evaluation of the quadrature rules, but I have to change the variable of the outer integration without having to rebuild the inner lazy evaluation tree.

[JordiManyer]

Hi @tam724 , thanks for the detailed issue.

Would it make sense to have cache invalidation? Yes, I think so. What is the best way of doing it? That, I am not sure about.
We obviously don't want to turn it off in general, and we don't want our solution to affect the regular code.

With that in mind, I would either

• do what you have done (in a general way that catches all possibilities). This however requires the user to go low-level like you did (which I think should be fine?)
• maybe add it as a preference that requires re-compilation of the code? Like a static if, basically. But then this would invalidate caches every time, all throughout the code.

I am slightly leaning towards the first one, to be honest. Seems less invasive.

[Antoinemarteau]

The issue only appears when you go low level:

Gridap.assemble_vector(b, V)[1] # uses α = 1.
α[] = 2.0
Gridap.assemble_vector(b, V)[1] # uses α = 2.

so I think it makes sense to have a low level solution anyway.

To my understanding, this LazyArray cache does two things:

• provide allocated memory (a cache) to maps for them to evaluate in place without allocation
• memoize the last array access to reuse the computed results

The issue here relates to memoization /the IndexItemPair part of the cache. So I see a third possibility, a kwarg to getindex! that tells to ignore all memoized values (no_memo? recompute_value? skip_cached_val?) , on the line of:

function getindex!(cache, a::LazyArray, i::Integer; no_memo::Val{B}) # B = true or false
  if B 
    return _no_memo_getindex!(cache, a, i)
  else
    return getindex!(cache,a,i)
  end
end

function _no_memo_getindex!(cache, a::LazyArray, i::Integer)
  cg, cgi, cf, index_and_item =_cache
  index_and_item.index = -1
  index = LinearIndices(a)[i]
  gi = getindex!(cg, a.maps, i; no_memo=Val(true))
  index_and_item.item = _getindex_and_call!(cgi,gi,cf,a.args,i; no_memo=Val(true))
  index_and_item.item
end

We could also have another array_cache that puts nothing in place of the IndexItemPair, and adapt getindex! accordingly. This would completely deactivate memoization in this cache, and I feel that can be efficiently done with dispatches only (no @static stuff).

[tam724]

Hey, thanks for the reply!
Sure, this is low-level and probably of small importance in typical use of Gridap.
But the version of having a array_cache(a::AbstractArray; memoize_last_index=true) and "let the cache know whether to memoize the last index" sounded clean to me.

This could be implemented as:

mutable struct IndexItemPair{T,V}
  memoize_last_index::Bool
  index::T
  item::V
end

function is_memoized(index_and_item::IndexItemPair{T}, index::T) where T
  return index_and_item.memoize_last_index && index_and_item.index == index
end

I implemented a draft that misappropriates the cache-hash-table dict (that is passed down anyways) to pass the memoization-preference through the array_cache recursion.
Let me know what you think!

[Antoinemarteau]

Hey @tam724, no worries.

Thanks for the draft, I reviewed it but have two issues with this approach:

• The dict might contain both caches with and without last access memoization, so I wouldn't pass the kwarg in there
• Adding memoize_last_index::Bool into IndexItemPair has both a memory and runtime performance overhead w.r.t. the original code, I would like better to replace the ::IndexItemPair with nothing in the cache to resolve the branch at compile time using:

_is_memoized(index_and_item::IndexItemPair, index) = index_and_item.index == index
_is_memoized(::Nothing, index) = false

I made a new draft here.

Still, this solution is different with just invalidating all IndexItemPairs of the caches once, because one might want to use memoization by default but invalidate it from time to time.
@JordiManyer do you have a prefered solution ? If we want a "no memoization" LazyArray cache, is my draft fine (I fear passing kwargs is a bit invasive) ?