Jitted function 20X slower if it returns solutions

[TRo202]

Hi.

I'm trying to solve multiple QPs by using vmap and JIT, but if the function runQP (where it solves QP) has the line 'return solutions' in the end, it takes 0.7s.

def runQP(...):
    ...
    sol = self.qp.run(params_obj=(Q, c), params_eq=A, params_ineq=(l, u)).params
    
    return sol.primal

 ...

auto_batch_runQP = jax.vmap(self.qp_solver.runQP)
jitted_runQP_vectorized = jax.jit(auto_batch_runQP)
optimal_sol = jitted_runQP_vectorized(...)

In contrast, if there isn't return sol.primal , then it takes about 0.03s. But then I can only have access to traced values.

def runQP(...):
    ...
    sol = self.qp.run(params_obj=(Q, c), params_eq=A, params_ineq=(l, u)).params
    
    self.qp_solutions = sol.primal

 ...

auto_batch_runQP = jax.vmap(self.qp_solver.runQP)
jitted_runQP_vectorized = jax.jit(auto_batch_runQP)
jitted_runQP_vectorized(...)

My questions are as follows:

1. Can anyone please help me explain why including 'return' part makes the computation much slower? Is this normal or have I made any mistake?

2. JIT uses traced values, so jitted_runQP_vectorized should return traced values, but it returns the actual data values. So does this mean that the function runQP is not jitted properly?

3. Is there any other way to make this computation much faster? My goal is to solve and return solutions within 20ms.

I will sincerely appreciate all the comments and suggestions!

[jakevdp]

JAX's JIT compilation assumes the function you pass it is pure (see JAX sharp bits: pure functions. Under this assumption, the compiler can elliminate any computation that doesn't affect the return value. When you jit-compile a function with no return value, it means all its contents can be elliminated: and running an empty function will generally be much faster than running a funciton that does some number of computations.

The root of the issue here is that you're JIT-compiling a function with impure semantics: that is, it appears to rely on side-affects like in-place mutation of the self argument. This will not work in the way you expect: you need to make sure you're only using jit, vmap, and other JAX transformations with pure functions.

Hope that helps!

[jakevdp]

I don't know what self.qp.run() is doing, but the code you've included here looks fine in terms of side-effects.

[TRo202]

Oh I totally forgot that in-place update is not allowed for jnp. Then if I want to make a code that solves QPs every control cycle with variable input data, should I first initialize the QP form and update the matrices and vectors of the QP by x.at[idx].set(y) ? Is this okay?

Conventionally, we first initialize QP once, then update matrices and vectors and solve QP for every control cycle. But since in-place update is not allowed for jnp, I assume it will be computationally too expensive to update values by x.at[idx].set(y) for all QPs.

Could you please tell me if there is any better method to update matrices and vectors?

[jakevdp]

You can definitely update the matrices and vectors every iteration, you just can't do it via a side effect: for example, you can include the updated arrays in the return value of the JIT-compiled function.

[TRo202]

Umm sorry I'm not quite sure if I understood correctly..

So if the JIT-compiled function returns the updated arrays, there would be no problem in updating matrices and vectors, right?

class QPSolver:
    def __init__(self):
        self.qp = BoxOSQP(tol=1e-3)

    def runQP(self, nom_values, min_values, max_values):
        # update values inside c, A, l, u jnp arrays
        sol = self.qp.run(params_obj=(Q, c), params_eq=A, params_ineq=(l, u)).params
        return sol.primal, c, A, l, u

[jakevdp]

As I mentioned above, there's no way for me to know what self.qp.run does, so there's no way to know whether this is OK under JIT.