need to vectorize efficiently calculating only certain values in the matrix multiplication A * B, using a logical array L the size of A * B

Question

I have matrices A (m by v) and B (v by n). I also have a logical matrix L (m by n).

I am interested in calculating only the values in A * B that correspond to logical values in L (values of 1s). Essentially I am interested in the quantity ( A * B ) .* L .

For my problem, a typical L matrix has less than 0.1% percent of its values as 1s; the vast majority of the values are 0s. Thus, it makes no sense for me to literally perform ( A * B ) .* L , it would actually be faster to loop over each row of A * B that I want to compute, but even that is inefficient.

---

Possible solution (need help vectorizing this code if possible)

My particular problem may have a nice solution given that the logical matrix L has a nice structure.

Here's an example of L for a very small scale example (in most applications L is much much bigger and has much fewer 1-yellow entries, and many more 0-blue entries).

This L matrix is nice in that it can be represented as something like a permuted block matrix. This L in particular is composed of 9 "blocks" of 1s, where each block of 1s has its own set of row and column indices, defining a particular submatrix in A * B. For instance, the highlighted area here can be seen the values of 1 as a particular submatrix in L.

My solution was to do this: break the problem into submatrices over these blocks, and do matrix multiplications over each submatrix. I can get the row indices and column indices per each block's submatrix in L, organized in two cell lists "rowidxs_list" and "colidxs_list", both with the number of cells equal to the number of blocks. For instance in the block example I gave, subblock 1, I could calculate those particular values in A * B by simply doing A( rowidxs_list{1} , : ) * B( : , colidxs_list{1} ) .

That means that if I precomputed rowidxs_list and colidxs_list (ignore the costs of calculating these lists, they are negligable for my application), then my problem of calculating C = ( A * B ) .* L could effectively be done by:

C = sparse( m,n )
for i = 1:length( rowidxs_list )
    C( rowidxs_list{i} , colidxs_list{i} ) = ...
        A( rowidxs_list{i} , : ) * B( : , colidxs_list{i} );
end

This seems like it would be the most efficient way to solve this problem if I knew how to vectorize this for loop. Does anyone see a way to vectorize this?

There may be ways to vectorize if certain things hold, e.g. only if rowidxs_list and colidxs_list are matrix arrays instead of cell lists of lists (where each column in an array is an index list, thus replacing use of rowidxs_list{i} with rowidxs_list(i,:) ). I'd prefer to use cell lists here if possible since different lists can have different numbers of elements.

---

other suggested solution (creating a mex file?)

I first posted this question on the /r/matlab subreddit, see here for the reddit thread. The user "qtac" recommended that a C-MEX function linking to C programming language:

My gut feeling is the only way to really optimize this is with a C-MEX solution; otherwise, you are going to get obliterated by overhead from subsref in these loops. With C you could loop over L until you find a nonzero element, and then do only the row-column dot product needed to populate that specific element. You will miss out on a lot of the BLAS optimizations but the computational savings may make up for it. Honestly I bet an LLM could write 90%+ of that MEX function for you; it's a well-formulated problem.

I think this could be a good solution to pursue, but I'd like other opinions as well.

Answer 1

The (i , j)-th entry of the result matrix C can be seen as the vector product, without conjugation, of the i-th row of A (transposed) and the j-th column of B. So you can get all pairs of indices (i , j) from L and do that in a vectorized way:

% Example data:
m = 3; n = 4; v = 9;
L = logical([0 0 1 0; 0 1 1 0; 1 0 0 0]); % m by n
A = rand(m, v);
B = rand(v, n);

% Computation:
[ii, jj] = find(L); % all pairs for which the result is needed
C = sparse(m, n); % preallocate as sparse
ind = ii + (jj-1)*m; % linear index from ii and jj
C(ind) = sum(A(ii,:).'.*B(:,jj), 1); % actual computation

If you only need a vector containing the relevant entries of C, instead of the whole matrix:

% Computation:
[ii, jj] = find(L); % all pairs for which the result is needed
ind = ii + (jj-1)*m; % linear index from ii and jj
c = sum(A(ii,:).'.*B(:,jj), 1); % actual computation