Use JIT to skfem.asm

[ModelAdminIFCT]

Hi, I have read some old discussions and some of them used JIT to speed up the code.
I also want to do something similar, but I'm not fully sure how to apply JIT to skfem. Especially, a code I'm working on right now does a lot of skfem.asm() to build jacobians iteratively, and that makes the speed of the code super slow.
Is there a simpler example that uses JIT with scikit-fem? such as the Poisson example?

Thank you, as always.

[kinnala]

There is no definite answers yet, any ideas are welcome. Some facts:

• If you have N basis functions per element, the form function (wrapped by BilinearForm) gets evaluated N^2 times
• The inputs u and v to the form are largeish Nelems x Nqp numpy arrays
• Operations such as u * v are pretty fast since they are done within numpy
• If you have lots of such operations, it involves going back and forth between numpy low level and Python high level code
• This will make complicated forms slow

To minimize the back and forth, you should try to do as much as possible using a single numpy call within the form function. Another idea is to JIT the entire function, to evaluate everything in one low level call. We are still bound to do N^2 evaluations of the JITted form.

Can you give examples of the forms you assemble and mesh size?

[kinnala]

You could try profiling the form by printing time at the start of the form and before return. Then you can compare N^2 * time_per_evaluation to the total assembly time to see how much you could theoretically save by making the form evaluation faster.

[Yeop-Jung]

Hi, I'm sorry, I didn't quite understand all your answers. I'll try to ask again where I couldn't understand in your answer. Also, I'll try to give you an example of my case (it became too big now so it's impossible to put the entire code). And Thank you for the answer btw!

Please correct me if I'm wrong.
From your answer:

> If you have N basis functions per element, the form function (wrapped by BilinearForm) gets evaluated N^2 times
This is equivalent to the N variables to solve, am I right? and the form function in Poisson example would be:

@BilinearForm
def laplace(u, v, _):
    return dot(grad(u), grad(v))

> The inputs u and v to the form are largeish Nelems x Nqp numpy arrays
elems and qp stands for N elements and N quadrature point?
And is this the reason for the N^2 times evaluations in the bilinearform?

> If you have lots of such operations, it involves going back and forth between numpy low level and Python high level code
> This will make complicated forms slow
I think this would be why my code is so slow. I use iterative methods and it rebuild the jacobian matrix again every iteration.

> To minimize the back and forth, you should try to do as much as possible using a single numpy call within the form function. Another idea is to JIT the entire function, to evaluate everything in one low level call. We are still bound to do N^2 evaluations of the JITted form.
What's the single numpy call within the form function? How that's different from JITting the entire function? Sorry, I'm a real novice here, I'm not sure what the low level and high level calls are.

My form function is like this:

@fem.BilinearForm
  def Jij(u, v, p):
      w1 = p['sp1']
      A =  p['cond']
      DA = p['dcond']
      Df = p['df']
      
      return inner(A*grad(u),grad(v)) + inner(DA*u*grad(w1), grad(v)) - inner(Df*u,v)

I have 8 coupled transport equations (Poisson-typed), and because they are coupled either in the diffusion coefficient (A) or f .
Also I solve this code multiple times. Meaning that, there is a line of segments and I solve the single problem per segment. The result of the first segment becomes the inputs of the next segment and this goes on until the last n segment.
So, the code solve the single case of this problem iteratively and this single case repeated n times.

So there are # of single case iterations * n times skfem.asm process.

For now, I'm using 1D mesh with 200 - 1000 mesh elements, but since there are so many iterations, it's already slow.
Ultimately, I want to expand the code to run in 2D, but with this speed, it seems it's not very viable or usable.

If you can direct me to a simpler (relatively) example, that would be really helpful.
Please let me know if you need something more to understand the question or my issue. I'll do my best to inform you what I want to do.

Thank you.

[Yeop-Jung]

I also saw your old discussion and one of your code was:

import numpy as np
import timeit
import skfem as fe
from numba import jit

@jit(nogil=True, nopython=True)
def nlaplace(out, du, dv):
    for i in range(du.shape[1]):
        for j in range(du.shape[2]):
            for k in range(du.shape[0]):
                out[i, j] += du[k, i, j] * dv[k, i, j]

start = timeit.default_timer()
mesh = fe.MeshQuad1.init_tensor(*[np.linspace(0, 1, 1001)]*2)
element = fe.ElementQuad1()
mapping = fe.MappingIsoparametric(mesh, element)
basis = fe.Basis(mesh, element, mapping, intorder=2)

@fe.BilinearForm(nthreads=4)
def laplace(u, v, w):
    out = np.empty_like(u.grad[0])
    nlaplace(out, u.grad, v.grad)
    return out


lap = laplace.assemble(basis)
stop = timeit.default_timer()
print(stop - start)

Should I make my form as you did here? Like, instead of using the dot function, make a function similar to nlaplace in the above example?
And here you used (nthreads=4) is this something available without any parallelization of the code?

[kinnala]

It is hard for me to see what kind of 1D problem would make the finite element assembly a bottleneck. I have solved some nonlinear 3D problems with much more degrees-of-freedom but in these cases the linear solve is usually dominating by a order of magnitude because the form is very, very complex. However, in 1D I can see that the linear solve is also very fast so the relative cost of assembly might be higher. The form you have written down is not particularly complex, however, and it won't explain the slowness. I'm afraid I cannot help further without a minimal example which I can run and profile on my own computer.

[Yeop-Jung]

You're right. I did some profiles and found that the linear solver is pretty fast for my 1D problem but it was when building jacobian matrices multiple times with complicated coefficients. I'll try to make a simplified example as soon as possible.
Thanks a ton!