I am trying to see how can I improve the computational performance of this small function that I am using in python. I am using numba with nopython mode and in general, I observe some small benefit in contrast to just using standard python/numpy operations. def _k_points_eval_s(zk,zk_p1,zkp1_zk,k1_denom_det,k2_denom_det,nk,nk_p1): @njit(cache=True,fastmath=True) def k_points_eval_in_s(m_b_array,b_array,pk,pk_p1,m_k,tk_nom,tk): ##Ok.. now we want to have here a loop to evaluate agains the different parts ###Let's try to make a single broadcast for m_b_array first to the shape of the comparing elements for i in prange(0,m_b_array.shape[0],1): m_zp1_diff = m_b_array- zk_p1 # k2 m_z_diff = m_b_array- zk ##k1 # Now, we can calculate the determinant terms for the numerators k1 = (m_z_diff[:, 0]*zkp1_zk[:, 1] -m_z_diff[:, 1]*zkp1_zk[:, 0]) / (k1_denom_det) k2 = (m_zp1_diff[:, 0]*zkp1_zk[:, 1] -m_zp1_diff[:, 1]*zkp1_zk[:, 0]) / (k2_denom_det) ##Following, we need to compute pk and pk_p1 pk[:,0] = zk[:,0] + k1*nk[:,0] pk[:,1] = zk[:,1] + k1*nk[:,1] pk_p1[:,0] = zk_p1[:,0] +k2*nk_p1 [:,0] pk_p1[:,1] = zk_p1[:,1] +k2*nk_p1[:,1] # We now need to compute the difference between pk and zk pk_zk_diff = pk - zk pkp1_zkp1_diff = pk_p1 - zk_p1 # now we need to evaluate the dk and dk_p1 conditions.. # For dk we need, y - zk with y taking the values of: pk_p1 and zk_p1 dk_y1 = pk_p1 - zk dk_y2 = zk_p1 - zk # For dk_p1 we need, y - zk_p1 with y taking the values of: pk and zk dk_p1_y1 = pk - zk_p1 dk_p1_y2 = zk - zk_p1 # Now we need to calculate the determinants from these components. # dk_pk_p1 dk_pk_p1 = pk_zk_diff[:, 0]*dk_y1[:, 1] - pk_zk_diff[:, 1]*dk_y1[:, 0] # dk_zk_p1 dk_zk_p1 = pk_zk_diff[:, 0]*dk_y2[:, 1] - pk_zk_diff[:, 1]*dk_y2[:, 0] # dk_p1_pk dk_p1_pk = pkp1_zkp1_diff[:, 0]*dk_p1_y1[:,1] - pkp1_zkp1_diff[:, 1]*dk_p1_y1[:, 0] # dk_p1_zk dk_p1_zk = pkp1_zkp1_diff[:, 0]*dk_p1_y2[:,1] - pkp1_zkp1_diff[:, 1]*dk_p1_y2[:, 0] # We now need to evaluate the projections on the line segments. We first compute the norm for the denominator shown in 6.18b # We calculate the product as in equations 6.17a and 6.17b cond_1 = dk_pk_p1*dk_zk_p1 cond_2 = dk_p1_pk*dk_p1_zk tk_nom = m_b_array - np.where(((cond_1 < 0) & (cond_2 < 0))[:, np.newaxis], pk_p1, pk) m_b_pk_diff = m_b_array- pk pk_diff = pk_p1 - pk pk_norm = np.sqrt(pk_diff[:, 0]**2 + pk_diff[:, 1]**2) dot_1 = m_b_pk_diff[:,0]*pk_diff[:,0] + m_b_pk_diff[:,1]*pk_diff[:,1] tk_nom_norm = (tk_nom[:, 0]**2 + tk_nom[:, 1]**2)**(1/2) tk[:] = np.where(((dot_1 >= 0.0) & (dot_1 <= pk_norm**2)),tk_nom_norm/pk_norm, 10**12) m_k[:,0] = (zk[:,0] + tk*(zkp1_zk[:, 0])) m_k[:,1] = (zk[:,1] + tk*(zkp1_zk[:,1])) m_m_k_norm = 0.5 *( ((m_b_array[i,:,0]-m_k[:,0])**2) + ( (m_b_array[i,:,1]-m_k[:,1])**2 ) ) arg_min = np.argmin(m_m_k_norm) ##This returns a single scalar value b_array[i,:] = m_k[int(arg_min),:] return b_array return k_points_eval_in_s In general this is what I am performing here: -I need to run this function over multiple iterations and at the moment, this function is dominating my overall algorithm execution (the rest of the algorithm consist of solving some linear systems for which I am using PETSc) -For this function, I first have some reference arrays of shape (M,2) and (M,) which are not changing over iteration. Therefore, I use these as closure variables and are the ones being provided on the first call to _k_points_eval_s. -Then on each iteration, I provide m_b_array which is a numpy array of shape (N,M,2) and b_array with shape (N,2). The first one is the one used to perform the needed evaluations and the later one is where I assign the computed results. The other variables passed in this case are just some numpy arrays which are already initialized so that I avoid allocating memory for some intermediate results on each iteration. -I have also checked that all my arrays are continuous in memory (C order) However, despite the fact that I am using numba in this case, this function is still being the bottle neck of my overall algorithm. I was trying to see how to squeeze more performance by using the debug options mentioned in https://numba.pydata.org/numba-doc/dev/user/faq.html via import llvmlite.binding as llvm llvm.set_option('', '--debug-only=loop-vectorize') -In general, the output from LLVM is not trivial and as you would imagine, vectorization is not possible throught out most of my function. Having said this, here my question would be then: Is there anything I could do to improve the running performance of this function? I was thinking of expanding all operations via jitted function loops but I am not really sure if it is really worth the extra implementation. What do you think? Would there be any significant benefit from expanding all operations? For that, would it work to define additional jitted numba functions which are then called from within this main function? From what I have read, numpy operations are already highly optimized and I am not sure if implementing these operations on my own would be really beneficial and allow for vectorization. Thanks a lot for your comments. Continue reading...
